From af7a3c66aff97886cbab45b5cea67417e6ae01e6 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Wed, 2 Sep 2026 10:16:57 +0000 Subject: [PATCH] perf: parallelize YCSB key preload Preload writes went one at a time: preloadKeys waited for each Put or PutStream before issuing the next, so the configured concurrency never applied to the load phase. One deterministic producer now generates keys and values in the original order (the RNG state after preload is unchanged), a worker pool sized by the configured concurrency performs the puts, and one collector owns the accounting, the 10% failure threshold, and the spinner. The job queue is bounded so values are released as their put completes. Rebased onto main after #260: the pool carries both the plain and the CAS preload namespaces and keeps the puts-based accounting; the tests were ported onto the CAS-aware test server. The PR's commits were squashed by the rebase. Signed-off-by: Perfloop Agent Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YCoN4Him5ZY4hZovVp37ra --- client/cmd/ycsb/ycsb.go | 171 +++++++++++++++-------- client/cmd/ycsb/ycsb_test.go | 256 ++++++++++++++++++++++++++++++++++- 2 files changed, 370 insertions(+), 57 deletions(-) diff --git a/client/cmd/ycsb/ycsb.go b/client/cmd/ycsb/ycsb.go index b94dd1e0..c5ee0b6d 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 1bdbefc6..77e98385 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)