diff --git a/client/cmd/ycsb/ycsb.go b/client/cmd/ycsb/ycsb.go index b94dd1e..c5ee0b6 100644 --- a/client/cmd/ycsb/ycsb.go +++ b/client/cmd/ycsb/ycsb.go @@ -12,7 +12,6 @@ import ( "slices" "strings" "sync" - "sync/atomic" "time" "github.com/pterm/pterm" @@ -209,26 +208,102 @@ func preloadKeys(ctx context.Context, cfg YCSBConfig, ws WorkloadSpec, rng *rand expectedPuts += cfg.NumKeys } - var preloadErrors int32 - var successCount int32 - errorCh := make(chan error, 100) + useStreaming := cfg.ForceStreaming || cfg.ValueSize > StreamingThreshold + workerCount := max(1, min(cfg.Concurrency, max(expectedPuts, 1))) + + // The preload runs the configured concurrency: one deterministic producer + // generates keys and values in the same order as before (so the RNG state + // after preload is unchanged), a bounded worker pool performs the puts, and + // one collector owns the accounting and spinner output. The job queue is + // bounded so values are released as their put completes instead of holding + // the whole corpus in memory. + type preloadJob struct { + seq int // generation order, for a stable "first error" + key string + value []byte + cas bool + } + type preloadResult struct { + seq int + key string + err error + } + jobs := make(chan preloadJob, workerCount) + results := make(chan preloadResult, workerCount) + + var workerWg sync.WaitGroup + workerWg.Add(workerCount) + for range workerCount { + go func() { + defer workerWg.Done() + for { + select { + case <-ctx.Done(): + return + case job, ok := <-jobs: + if !ok { + return + } + if ctx.Err() != nil { + return + } + // Plain read/update keys and CAS keys live in disjoint namespaces + // so a key is never both plain-written and CAS-guarded (an + // unsupported mix). CAS keys are created with put-if-absent; a + // mismatch there means the key survived from an earlier run and + // is already CAS-owned, which is fine. + var err error + switch { + case job.cas && useStreaming: + _, err = client.PutStreamIfVersion(ctx, job.key, bytes.NewReader(job.value), 0, 0) + case job.cas: + _, err = client.PutIfVersion(ctx, job.key, job.value, 0, 0) + case useStreaming: + err = client.PutStream(ctx, job.key, bytes.NewReader(job.value), 0) + default: + err = client.Put(ctx, job.key, job.value, 0) + } + if outcome, _ := ClassifyCASResult(err); outcome == CASMismatch { + err = nil // already present from an earlier run: fine + } + results <- preloadResult{seq: job.seq, key: job.key, err: err} + } + } + }() + } - for i := range cfg.NumKeys { - // Check for context cancellation - select { - case <-ctx.Done(): - if spinner != nil { - spinner.Warning(fmt.Sprintf("Preload cancelled after %d/%d keys", i, cfg.NumKeys)) + // Completion order is independent of generation order, so keep the + // lowest-sequence error for the same first-error message the serial loop + // produced. + var collectWg sync.WaitGroup + collectWg.Add(1) + successCount, totalErrors, completedCount := 0, 0, 0 + firstErrorSeq := -1 + var firstError error + go func() { + defer collectWg.Done() + for result := range results { + completedCount++ + if result.err != nil { + totalErrors++ + if firstErrorSeq < 0 || result.seq < firstErrorSeq { + firstErrorSeq = result.seq + firstError = fmt.Errorf("key %s: %w", result.key, result.err) + } + } else { + successCount++ + } + if spinner != nil && (completedCount == 1 || completedCount%100 == 0) { + spinner.UpdateText(fmt.Sprintf("Preloading keys: %d/%d (errors: %d)", + completedCount, expectedPuts, totalErrors)) } - return ctx.Err() - default: } + }() - useStreaming := cfg.ForceStreaming || cfg.ValueSize > StreamingThreshold - // Plain read/update keys and CAS keys live in disjoint namespaces so a key - // is never both plain-written and CAS-guarded (an unsupported mix). CAS - // keys are created with put-if-absent; a mismatch there means the key - // survived from an earlier run and is already CAS-owned, which is fine. + cancelled := false + seq := 0 +produce: + for i := range cfg.NumKeys { for _, p := range []struct { enabled bool key string @@ -240,58 +315,42 @@ func preloadKeys(ctx context.Context, cfg YCSBConfig, ws WorkloadSpec, rng *rand if !p.enabled { continue } - val := generateValue(rng, cfg.ValueSize) - var err error - switch { - case p.cas && useStreaming: - _, err = client.PutStreamIfVersion(ctx, p.key, bytes.NewReader(val), 0, 0) - case p.cas: - _, err = client.PutIfVersion(ctx, p.key, val, 0, 0) - case useStreaming: - err = client.PutStream(ctx, p.key, bytes.NewReader(val), 0) - default: - err = client.Put(ctx, p.key, val, 0) - } - if outcome, _ := ClassifyCASResult(err); outcome == CASMismatch { - err = nil // already present from an earlier run: fine + // Check for cancellation before generating the next owned value. + if ctx.Err() != nil { + cancelled = true + break produce } - - if err != nil { - atomic.AddInt32(&preloadErrors, 1) - select { - case errorCh <- fmt.Errorf("key %s: %w", p.key, err): - default: // Don't block on error channel - } - } else { - atomic.AddInt32(&successCount, 1) + val := generateValue(rng, cfg.ValueSize) + select { + case <-ctx.Done(): + cancelled = true + break produce + case jobs <- preloadJob{seq: seq, key: p.key, value: val, cas: p.cas}: + seq++ } } - - if spinner != nil && i%100 == 0 { - spinner.UpdateText(fmt.Sprintf("Preloading keys: %d/%d (errors: %d)", - i+1, cfg.NumKeys, atomic.LoadInt32(&preloadErrors))) - } } - close(errorCh) + close(jobs) + workerWg.Wait() + close(results) + collectWg.Wait() - // Collect sample of errors - var sampleErrors []error - for err := range errorCh { - if len(sampleErrors) < 5 { - sampleErrors = append(sampleErrors, err) + if cancelled || ctx.Err() != nil { + if spinner != nil { + spinner.Warning(fmt.Sprintf("Preload cancelled after %d/%d puts", completedCount, expectedPuts)) } + return ctx.Err() } - totalErrors := atomic.LoadInt32(&preloadErrors) if totalErrors > 0 { if spinner != nil { spinner.Warning(fmt.Sprintf("Preloaded %d/%d puts (%d errors)", - atomic.LoadInt32(&successCount), expectedPuts, totalErrors)) + successCount, expectedPuts, totalErrors)) } - if int(totalErrors) > expectedPuts/10 { // If more than 10% failed, consider it a failure - if len(sampleErrors) > 0 { - return fmt.Errorf("preload failed with %d errors, first error: %w", totalErrors, sampleErrors[0]) + if totalErrors > expectedPuts/10 { // If more than 10% failed, consider it a failure + if firstError != nil { + return fmt.Errorf("preload failed with %d errors, first error: %w", totalErrors, firstError) } } } else { diff --git a/client/cmd/ycsb/ycsb_test.go b/client/cmd/ycsb/ycsb_test.go index 1bdbefc..77e9838 100644 --- a/client/cmd/ycsb/ycsb_test.go +++ b/client/cmd/ycsb/ycsb_test.go @@ -4,11 +4,16 @@ package ycsb import ( + "bytes" "context" + "errors" + "math/rand" "net" + "strings" "sync" "sync/atomic" "testing" + "time" "github.com/pterm/pterm" cacheclient "github.com/tigrisdata/ocache/client" @@ -30,6 +35,7 @@ type ycsbReadServer struct { mu sync.RWMutex values map[string][]byte versions map[string]uint64 // CAS version per key; 0 = absent + putDelay time.Duration // simulated per-put latency, honouring ctx getCalls atomic.Int64 responseBytes atomic.Int64 casBumps atomic.Int64 // successful PutObjectIfVersion applications @@ -39,7 +45,17 @@ func newYCSBReadServer() *ycsbReadServer { return &ycsbReadServer{values: make(map[string][]byte), versions: make(map[string]uint64)} } -func (s *ycsbReadServer) PutObject(_ context.Context, req *pb.PutRequest) (*pb.PutResponse, error) { +func (s *ycsbReadServer) PutObject(ctx context.Context, req *pb.PutRequest) (*pb.PutResponse, error) { + if s.putDelay > 0 { + timer := time.NewTimer(s.putDelay) + defer timer.Stop() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + } + } + data := append([]byte(nil), req.Data...) s.mu.Lock() @@ -133,6 +149,66 @@ func startYCSBReadServer(tb testing.TB) (*ycsbReadServer, string) { return cacheServer, listener.Addr().String() } +// preloadTestServer wraps ycsbReadServer to inject per-key put failures and +// observe put concurrency for the preload tests. +type preloadTestServer struct { + *ycsbReadServer + + putFailures map[string]error + putCalls atomic.Int64 + activePuts atomic.Int64 + maxPutConcurrency atomic.Int64 +} + +func (s *preloadTestServer) PutObject(ctx context.Context, req *pb.PutRequest) (*pb.PutResponse, error) { + s.putCalls.Add(1) + active := s.activePuts.Add(1) + defer s.activePuts.Add(-1) + for { + maxActive := s.maxPutConcurrency.Load() + if active <= maxActive || s.maxPutConcurrency.CompareAndSwap(maxActive, active) { + break + } + } + + if err, ok := s.putFailures[req.Key]; ok { + return nil, err + } + return s.ycsbReadServer.PutObject(ctx, req) +} + +func startPreloadTestServer(tb testing.TB, putDelay time.Duration, putFailures map[string]error) (*preloadTestServer, string) { + tb.Helper() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + tb.Fatal(err) + } + + grpcServer := grpc.NewServer() + cacheServer := &preloadTestServer{ + ycsbReadServer: newYCSBReadServer(), + putFailures: putFailures, + } + cacheServer.ycsbReadServer.putDelay = putDelay + pb.RegisterCacheServiceServer(grpcServer, cacheServer) + go func() { + _ = grpcServer.Serve(listener) + }() + + tb.Cleanup(func() { + grpcServer.Stop() + _ = listener.Close() + }) + + return cacheServer, listener.Addr().String() +} + +// plainWorkload is a read-only spec: the preload writes the plain namespace only. +func plainWorkload() WorkloadSpec { + return WorkloadSpec{Weights: [OpNum]float64{OpRead: 1}} +} + func readOnlyYCSBConfig(addr string, valueSize int) YCSBConfig { return YCSBConfig{ Addr: addr, @@ -159,6 +235,184 @@ func disablePtermOutput(tb testing.TB) { }) } +func BenchmarkPreloadYCSBKeys(b *testing.B) { + disablePtermOutput(b) + cacheServer, addr := startYCSBReadServer(b) + cacheServer.putDelay = time.Millisecond + cfg := YCSBConfig{ + Addr: addr, + ConnMode: string(cacheclient.ModeSimple), + ConnectionPoolSize: 8, + NumKeys: 64, + ValueSize: 100, + NumOps: 64, + Concurrency: 8, + Workload: "C", + Seed: 1, + NoProgress: true, + } + + for b.Loop() { + result, err := RunYCSBWithContext(context.Background(), cfg) + if err != nil { + b.Fatal(err) + } + if result.Errors != 0 { + b.Fatalf("RunYCSBWithContext reported %d errors", result.Errors) + } + } +} + +// TestPreloadKeysPreservesValuesAndRNG: the parallel preload writes the same +// key/value mapping as the serial one, leaves the RNG in the same state, and +// actually runs puts concurrently, bounded by the configured concurrency. +func TestPreloadKeysPreservesValuesAndRNG(t *testing.T) { + disablePtermOutput(t) + const ( + seed = int64(42) + numKeys = 16 + valueSize = 32 + concurrency = 4 + ) + cacheServer, addr := startPreloadTestServer(t, 5*time.Millisecond, nil) + cfg := YCSBConfig{ + Addr: addr, + ConnMode: string(cacheclient.ModeSimple), + ConnectionPoolSize: concurrency, + NumKeys: numKeys, + ValueSize: valueSize, + Concurrency: concurrency, + NoProgress: true, + } + + wantRNG := rand.New(rand.NewSource(seed)) + wantValues := make(map[string][]byte, numKeys) + for i := range numKeys { + wantValues[hashKey(i)] = generateValue(wantRNG, valueSize) + } + wantNextSeed := wantRNG.Int63() + + gotRNG := rand.New(rand.NewSource(seed)) + if err := preloadKeys(context.Background(), cfg, plainWorkload(), gotRNG); err != nil { + t.Fatal(err) + } + + if got := cacheServer.maxPutConcurrency.Load(); got < 2 { + t.Fatalf("maximum concurrent puts = %d, want at least 2", got) + } + if got := cacheServer.maxPutConcurrency.Load(); got > concurrency { + t.Fatalf("maximum concurrent puts = %d, want at most %d", got, concurrency) + } + cacheServer.mu.RLock() + defer cacheServer.mu.RUnlock() + if got, want := len(cacheServer.values), numKeys; got != want { + t.Fatalf("preloaded values = %d, want %d", got, want) + } + for key, want := range wantValues { + if got := cacheServer.values[key]; !bytes.Equal(got, want) { + t.Errorf("value for %q changed", key) + } + } + if got := gotRNG.Int63(); got != wantNextSeed { + t.Errorf("RNG state after preload = %d, want %d", got, wantNextSeed) + } +} + +// TestPreloadKeysPreservesErrorThreshold: the 10% failure threshold and the +// "first error" message are unchanged by parallel completion order. +func TestPreloadKeysPreservesErrorThreshold(t *testing.T) { + disablePtermOutput(t) + for _, tc := range []struct { + name string + failedKeys []int + wantFailure bool + }{ + {name: "at threshold", failedKeys: []int{0}, wantFailure: false}, + {name: "over threshold", failedKeys: []int{0, 1}, wantFailure: true}, + } { + t.Run(tc.name, func(t *testing.T) { + failures := make(map[string]error, len(tc.failedKeys)) + for _, keyNum := range tc.failedKeys { + failures[hashKey(keyNum)] = errors.New("injected preload failure") + } + cacheServer, addr := startPreloadTestServer(t, 0, failures) + cfg := YCSBConfig{ + Addr: addr, + ConnMode: string(cacheclient.ModeSimple), + ConnectionPoolSize: 4, + NumKeys: 10, + ValueSize: 8, + Concurrency: 4, + NoProgress: true, + } + + err := preloadKeys(context.Background(), cfg, plainWorkload(), rand.New(rand.NewSource(1))) + if tc.wantFailure { + if err == nil { + t.Fatal("preload succeeded, want threshold error") + } + if !strings.Contains(err.Error(), "key "+hashKey(tc.failedKeys[0])) { + t.Errorf("error = %v, want first failed key %q", err, hashKey(tc.failedKeys[0])) + } + } else if err != nil { + t.Fatalf("preload failed at threshold: %v", err) + } + if got := cacheServer.putCalls.Load(); got != int64(cfg.NumKeys) { + t.Errorf("put calls = %d, want %d", got, cfg.NumKeys) + } + }) + } +} + +// TestPreloadKeysWaitsForWorkersOnCancellation: cancelling the preload returns +// context.Canceled only after every in-flight put has finished, and no put is +// started afterwards. +func TestPreloadKeysWaitsForWorkersOnCancellation(t *testing.T) { + disablePtermOutput(t) + cacheServer, addr := startPreloadTestServer(t, 100*time.Millisecond, nil) + cfg := YCSBConfig{ + Addr: addr, + ConnMode: string(cacheclient.ModeSimple), + ConnectionPoolSize: 4, + NumKeys: 100, + ValueSize: 8, + Concurrency: 4, + NoProgress: true, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + resultCh := make(chan error, 1) + go func() { + resultCh <- preloadKeys(ctx, cfg, plainWorkload(), rand.New(rand.NewSource(1))) + }() + + deadline := time.Now().Add(time.Second) + for cacheServer.activePuts.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if cacheServer.activePuts.Load() == 0 { + t.Fatal("preload did not start a put") + } + cancel() + + select { + case err := <-resultCh: + if !errors.Is(err, context.Canceled) { + t.Fatalf("preload error = %v, want context.Canceled", err) + } + case <-time.After(time.Second): + t.Fatal("preload did not return after cancellation") + } + if got := cacheServer.activePuts.Load(); got != 0 { + t.Fatalf("active puts after preload returned = %d, want 0", got) + } + calls := cacheServer.putCalls.Load() + time.Sleep(20 * time.Millisecond) + if got := cacheServer.putCalls.Load(); got != calls { + t.Fatalf("put calls continued after preload returned: started at %d, now %d", calls, got) + } +} + func TestRunYCSBReadOnlyDrainsResponses(t *testing.T) { disablePtermOutput(t) cacheServer, addr := startYCSBReadServer(t)