diff --git a/Makefile b/Makefile index 06daf350..fa5141a5 100644 --- a/Makefile +++ b/Makefile @@ -282,7 +282,15 @@ build-bench-cache-service-delete: proto build-bench-cache-service-list-with-values: proto @test -n "$(PERFLOOP_BUILD_OUTPUT_DIR)" || { echo "PERFLOOP_BUILD_OUTPUT_DIR is required"; exit 1; } @mkdir -p "$(PERFLOOP_BUILD_OUTPUT_DIR)" - @cd server && CGO_ENABLED=1 CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" go test $(LDFLAGS) -c -o "$(PERFLOOP_BUILD_OUTPUT_DIR)/cache-service-list-with-values.test" ./service + @cd server && CGO_ENABLED=1 CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" go test $(LDFLAGS) -tags=ocache_benchmark -c -o "$(PERFLOOP_BUILD_OUTPUT_DIR)/cache-service-list-with-values.test" ./service + +.PHONY: test-bench-cache-service-list-with-values-cancellation +test-bench-cache-service-list-with-values-cancellation: proto + @cd server && CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" go test $(LDFLAGS) -tags=ocache_benchmark -v -timeout 60s -run '^TestCacheServiceListWithValuesCancellationReleasesBlockedRead$$' ./service + +.PHONY: test-race-bench-cache-service-list-with-values-cancellation +test-race-bench-cache-service-list-with-values-cancellation: proto + @cd server && CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" go test $(LDFLAGS) -race -tags=ocache_benchmark -v -timeout 60s -run '^TestCacheServiceListWithValuesCancellationReleasesBlockedRead$$' ./service .PHONY: run-background run-background: diff --git a/coordinator/benchmark.go b/coordinator/benchmark.go new file mode 100644 index 00000000..a248c274 --- /dev/null +++ b/coordinator/benchmark.go @@ -0,0 +1,29 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build ocache_benchmark + +package coordinator + +import ( + "github.com/tigrisdata/ocache/coordinator/ring" +) + +// NewBenchmarkCoordinator wires a running in-memory ring to the normal routing +// path without starting memberlist. It is used only by the cluster list +// cancellation benchmark. +func NewBenchmarkCoordinator(ringManager *ring.RingManager, localID string) *Coordinator { + return &Coordinator{ + config: &Config{ + Enabled: true, + MyNodeID: localID, + LifecyclerConfig: ring.LifecyclerConfig{ + RingConfig: ring.Config{ReplicationFactor: 1}, + }, + }, + ringManager: ringManager, + router: NewRouter(ringManager, localID), + stopCh: make(chan struct{}), + errCh: make(chan error, 1), + } +} diff --git a/coordinator/ring/topology_benchmark.go b/coordinator/ring/topology_benchmark.go index a50db9ac..4a9235aa 100644 --- a/coordinator/ring/topology_benchmark.go +++ b/coordinator/ring/topology_benchmark.go @@ -1,7 +1,7 @@ // Copyright 2026 Tigris Data, Inc. // SPDX-License-Identifier: Apache-2.0 -//go:build ocache_topology_benchmark +//go:build ocache_benchmark || ocache_topology_benchmark package ring diff --git a/server/operations/list_with_values.go b/server/operations/list_with_values.go index bedde0ae..a7ab66e8 100644 --- a/server/operations/list_with_values.go +++ b/server/operations/list_with_values.go @@ -60,18 +60,29 @@ func (o *Operations) ListLocalWithValues(ctx context.Context, prefix string, lim var lastKey string var hasMore bool + if err := ctx.Err(); err != nil { + return nil, "", false, err + } + err := retry.Do(ctx, retry.DefaultConfig(), "ListKeyValuesWithPagination", func() error { var listErr error - storageEntries, lastKey, hasMore, listErr = o.storage.ListKeyValuesWithPagination(prefix, startKey, limit) + storageEntries, lastKey, hasMore, listErr = o.storage.ListKeyValuesWithPagination(ctx, prefix, startKey, limit) return listErr }) if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, "", false, ctxErr + } return nil, "", false, err } - // Convert storage entries to proto entries + // Convert storage entries to proto entries. A canceled request never + // publishes a partially converted page. entries := make([]*pb.KeyValue, len(storageEntries)) for i, e := range storageEntries { + if err := ctx.Err(); err != nil { + return nil, "", false, err + } entries[i] = &pb.KeyValue{ Key: e.Key, Value: e.Value, diff --git a/server/service/list_with_values.go b/server/service/list_with_values.go index 1b710a16..23a21e70 100644 --- a/server/service/list_with_values.go +++ b/server/service/list_with_values.go @@ -12,11 +12,15 @@ import ( zlog "github.com/rs/zerolog/log" "github.com/tigrisdata/ocache/common/metrics" pb "github.com/tigrisdata/ocache/proto" + "github.com/tigrisdata/ocache/storage/benchio" ) // ListWithValues implements the public ListWithValues RPC. // Returns sorted, paginated key-value pairs with continuation token. func (s *CacheService) ListWithValues(ctx context.Context, req *pb.ListRequest) (*pb.ListWithValuesResponse, error) { + endBenchmarkHandler := benchio.BeginListHandlerForBenchmark() + defer endBenchmarkHandler() + start := time.Now() defer func() { metrics.RPCDuration.WithLabelValues("ListWithValues").Observe(float64(time.Since(start).Milliseconds())) @@ -52,6 +56,9 @@ func (s *CacheService) ListWithValues(ctx context.Context, req *pb.ListRequest) // ListLocalWithValues implements the internal node-local ListWithValues RPC. // Returns sorted, paginated key-value pairs from this node's local storage. func (s *CacheService) ListLocalWithValues(ctx context.Context, req *pb.ListRequest) (*pb.ListWithValuesResponse, error) { + endBenchmarkHandler := benchio.BeginListHandlerForBenchmark() + defer endBenchmarkHandler() + start := time.Now() defer func() { metrics.RPCDuration.WithLabelValues("ListLocalWithValues").Observe(float64(time.Since(start).Milliseconds())) diff --git a/server/service/list_with_values_cancellation_benchmark_test.go b/server/service/list_with_values_cancellation_benchmark_test.go new file mode 100644 index 00000000..6182a2fa --- /dev/null +++ b/server/service/list_with_values_cancellation_benchmark_test.go @@ -0,0 +1,359 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build ocache_benchmark && !ocache_topology_benchmark + +package service + +import ( + "bytes" + "context" + "fmt" + "net" + "testing" + "time" + + "github.com/stretchr/testify/require" + "github.com/tigrisdata/ocache/coordinator" + "github.com/tigrisdata/ocache/coordinator/ring" + pb "github.com/tigrisdata/ocache/proto" + stor "github.com/tigrisdata/ocache/storage" + "github.com/tigrisdata/ocache/storage/benchio" + "google.golang.org/grpc" +) + +const ( + canceledListBenchmarkCount = 1000 + canceledListBenchmarkValueSize = 128 * 1024 +) + +type canceledListBenchmarkEnvironment struct { + storage *stor.Storage + peerStorage *stor.Storage + service *CacheService + peerServer *grpc.Server + peerListener net.Listener + closeRing func() + prefix string +} + +func newCanceledListBenchmarkEnvironment(tb testing.TB, count, valueSize int) *canceledListBenchmarkEnvironment { + tb.Helper() + + newStorage := func() *stor.Storage { + storage, err := stor.NewStorageWithConfig(&stor.StorageConfig{ + DiskPath: tb.TempDir(), + InlineThreshold: 1, + CompactThreshold: 4 * 1024, + SegmentSize: 256 * 1024 * 1024, + CleanupInterval: time.Hour, + DisableRecompaction: true, + }) + require.NoError(tb, err) + return storage + } + + storage := newStorage() + peerStorage := newStorage() + ringManager, closeRing, err := ring.NewTopologyBenchmarkManager(2, 128) + require.NoError(tb, err) + coord := coordinator.NewBenchmarkCoordinator(ringManager, "member-0") + + peerListener, err := net.Listen("tcp", "127.0.0.1:10001") + require.NoError(tb, err) + peerServer := grpc.NewServer( + grpc.MaxRecvMsgSize(128*1024*1024), + grpc.MaxSendMsgSize(128*1024*1024), + ) + pb.RegisterCacheServiceServer(peerServer, NewCacheService(nil, peerStorage)) + go func() { _ = peerServer.Serve(peerListener) }() + + env := &canceledListBenchmarkEnvironment{ + storage: storage, + peerStorage: peerStorage, + service: NewCacheService(coord, storage), + peerServer: peerServer, + peerListener: peerListener, + closeRing: closeRing, + prefix: "list-values-cancel-benchmark-", + } + value := bytes.Repeat([]byte("x"), valueSize) + for i := 0; i < count; i++ { + for _, nodeStorage := range []*stor.Storage{storage, peerStorage} { + key := fmt.Sprintf("%s%04d", env.prefix, i) + if nodeStorage == peerStorage { + key = fmt.Sprintf("%speer-%04d", env.prefix, i) + } else { + key = fmt.Sprintf("%slocal-%04d", env.prefix, i) + } + require.NoError(tb, nodeStorage.Put(key, bytes.NewReader(value), 0)) + } + } + return env +} + +func (env *canceledListBenchmarkEnvironment) close() { + if env.peerServer != nil { + // Wait for in-flight peer handlers before closing their storage. The + // baseline intentionally holds a read gate, so Stop could return while + // its handler was still unwinding and let RocksDB close under its iterator. + env.peerServer.GracefulStop() + env.peerServer = nil + } + if env.peerListener != nil { + _ = env.peerListener.Close() + env.peerListener = nil + } + if env.storage != nil { + env.storage.Close() + env.storage = nil + } + if env.peerStorage != nil { + env.peerStorage.Close() + env.peerStorage = nil + } + if env.closeRing != nil { + env.closeRing() + env.closeRing = nil + } +} + +func TestCacheServiceListWithValuesCancellationReleasesBlockedRead(t *testing.T) { + quietCacheServiceBenchmarkLogs(t) + env := newCanceledListBenchmarkEnvironment(t, 2, canceledListBenchmarkValueSize) + defer env.close() + + started, release, restore := benchio.BlockPayloadReadsForBenchmark() + defer func() { + release() + restore() + }() + + type result struct { + response *pb.ListWithValuesResponse + err error + } + resultCh := make(chan result, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + response, err := env.service.ListWithValues(ctx, &pb.ListRequest{ + Prefix: env.prefix, + Limit: 2, + }) + resultCh <- result{response: response, err: err} + }() + + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("ListWithValues did not reach the payload read") + } + cancel() + + select { + case result := <-resultCh: + require.ErrorIs(t, result.err, context.Canceled) + require.Nil(t, result.response) + case <-time.After(2 * time.Second): + t.Fatal("ListWithValues remained blocked after cancellation") + } +} + +type canceledListBenchmarkStats struct { + cancelToExitNanos int64 + rowsAfterCancel int64 + bytesAfterCancel int64 + activeScansAfterCancel int64 + activeHandlersAfterCancel int64 + activeReadersAfterCancel int64 + handlerExitedBeforeGate int64 + operations int64 +} + +// waitForCanceledListBenchmarkFanout waits until both cluster branches have +// reached their first payload reader before cancellation is recorded. +func waitForCanceledListBenchmarkFanout() bool { + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if benchio.ActiveListScansForBenchmark() >= 2 && + benchio.ActiveListHandlersForBenchmark() >= 2 && + benchio.ActivePayloadReadersForBenchmark() >= 2 { + return true + } + time.Sleep(100 * time.Microsecond) + } + return false +} + +// BenchmarkCacheServiceListWithValuesCancellation measures the public list +// operation after its first raw-file payload read has started. The benchmark +// records cancellation delay, post-cancel scan work, and ownership state while +// a gate holds the first payload read. +func BenchmarkCacheServiceListWithValuesCancellation(b *testing.B) { + quietCacheServiceBenchmarkLogs(b) + env := newCanceledListBenchmarkEnvironment(b, canceledListBenchmarkCount, canceledListBenchmarkValueSize) + defer env.close() + + req := &pb.ListRequest{Prefix: env.prefix, Limit: canceledListBenchmarkCount} + var stats canceledListBenchmarkStats + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + b.StopTimer() + stats.operations++ + benchio.ResetPayloadStatsForBenchmark() + started, release, restore := benchio.BlockPayloadReadsForBenchmark() + ctx, cancel := context.WithCancel(context.Background()) + ready := make(chan struct{}) + cancelled := make(chan time.Time, 1) + fanoutReady := make(chan bool, 1) + activeStats := make(chan canceledListBenchmarkStats, 1) + postCancelStats := make(chan canceledListBenchmarkStats, 1) + handlerDone := make(chan struct{}) + released := make(chan struct{}) + go func() { + close(ready) + <-started + fanoutReady <- waitForCanceledListBenchmarkFanout() + cancelAt := time.Now() + benchio.MarkPayloadCancellationForBenchmark() + cancel() + cancelled <- cancelAt + time.Sleep(5 * time.Millisecond) + activeStats <- canceledListBenchmarkStats{ + activeScansAfterCancel: benchio.ActiveListScansForBenchmark(), + activeHandlersAfterCancel: benchio.ActiveListHandlersForBenchmark(), + activeReadersAfterCancel: benchio.ActivePayloadReadersForBenchmark(), + } + time.Sleep(5 * time.Millisecond) + release() + close(released) + <-handlerDone + postCancelStats <- canceledListBenchmarkStats{ + rowsAfterCancel: benchio.PostCancellationRowsForBenchmark(), + bytesAfterCancel: benchio.PostCancellationBytesForBenchmark(), + } + }() + <-ready + + b.StartTimer() + _, _ = env.service.ListWithValues(ctx, req) + b.StopTimer() + close(handlerDone) + + cancelAt := <-cancelled + stats.cancelToExitNanos += time.Since(cancelAt).Nanoseconds() + fanoutReached := <-fanoutReady + select { + case <-released: + default: + stats.handlerExitedBeforeGate++ + } + active := <-activeStats + stats.activeScansAfterCancel += active.activeScansAfterCancel + stats.activeHandlersAfterCancel += active.activeHandlersAfterCancel + stats.activeReadersAfterCancel += active.activeReadersAfterCancel + <-released + postCancel := <-postCancelStats + stats.rowsAfterCancel += postCancel.rowsAfterCancel + stats.bytesAfterCancel += postCancel.bytesAfterCancel + restore() + cancel() + if !fanoutReached { + b.Fatal("cluster cancellation workload did not reach both handlers, scans, and payload readers") + } + b.StartTimer() + } + b.StopTimer() + + operations := float64(stats.operations) + b.ReportMetric(float64(stats.cancelToExitNanos)/operations, "cancel-to-exit-ns/op") + b.ReportMetric(float64(stats.rowsAfterCancel)/operations, "rows-after-cancel/op") + b.ReportMetric(float64(stats.bytesAfterCancel)/operations, "bytes-after-cancel/op") + b.ReportMetric(float64(stats.activeScansAfterCancel)/operations, "active-list-scans-after-cancel/op") + b.ReportMetric(float64(stats.activeHandlersAfterCancel)/operations, "active-list-handlers-after-cancel/op") + b.ReportMetric(float64(stats.activeReadersAfterCancel)/operations, "active-payload-readers-after-cancel/op") + b.ReportMetric(float64(stats.handlerExitedBeforeGate)/operations, "handler-exit-before-gate/op") +} + +func consumeClusterListWithValuesBenchmarkResponse(b *testing.B, response *pb.ListWithValuesResponse, count, valueSize int, expectMore bool) []string { + b.Helper() + if response == nil { + b.Fatal("ListWithValues returned a nil response") + } + if len(response.Entries) != count { + b.Fatalf("ListWithValues returned %d entries, want %d", len(response.Entries), count) + } + if response.HasMore != expectMore { + b.Fatalf("ListWithValues returned unexpected has_more: got=%v want=%v", response.HasMore, expectMore) + } + if (response.ContinuationToken != "") != expectMore { + b.Fatalf("ListWithValues returned unexpected continuation: has_more=%v token=%q", response.HasMore, response.ContinuationToken) + } + + keys := make([]string, 0, len(response.Entries)) + var totalBytes int + previousKey := "" + for _, entry := range response.Entries { + if entry == nil { + b.Fatal("ListWithValues returned a nil entry") + } + if previousKey != "" && entry.Key <= previousKey { + b.Fatalf("ListWithValues returned keys out of order: %q after %q", entry.Key, previousKey) + } + previousKey = entry.Key + keys = append(keys, entry.Key) + if entry.ValueOmitted || entry.ValueLength != int64(valueSize) || len(entry.Value) != valueSize { + b.Fatalf("ListWithValues returned an invalid value: key=%q omitted=%v length=%d bytes=%d", entry.Key, entry.ValueOmitted, entry.ValueLength, len(entry.Value)) + } + totalBytes += len(entry.Value) + } + if totalBytes != count*valueSize { + b.Fatalf("ListWithValues returned %d value bytes, want %d", totalBytes, count*valueSize) + } + return keys +} + +// BenchmarkCacheServiceListWithValuesRawLive is the live-context guard for +// the same spilled-value consumer path. It checks that opening a private +// interruptible descriptor does not silently transfer material cost to +// successful raw-file pages while preserving cluster ordering and cursors. +func BenchmarkCacheServiceListWithValuesRawLive(b *testing.B) { + quietCacheServiceBenchmarkLogs(b) + env := newCanceledListBenchmarkEnvironment(b, 100, canceledListBenchmarkValueSize) + defer env.close() + + req := &pb.ListRequest{Prefix: env.prefix, Limit: 100} + b.ReportAllocs() + for b.Loop() { + ctx, cancel := context.WithCancel(context.Background()) + response, err := env.service.ListWithValues(ctx, req) + if err != nil { + b.Fatal(err) + } + firstKeys := consumeClusterListWithValuesBenchmarkResponse(b, response, 100, canceledListBenchmarkValueSize, true) + response, err = env.service.ListWithValues(ctx, &pb.ListRequest{ + Prefix: env.prefix, + Limit: 100, + ContinuationToken: response.ContinuationToken, + }) + cancel() + if err != nil { + b.Fatal(err) + } + secondKeys := consumeClusterListWithValuesBenchmarkResponse(b, response, 100, canceledListBenchmarkValueSize, false) + if firstKeys[len(firstKeys)-1] >= secondKeys[0] { + b.Fatalf("cluster continuation moved backwards: first=%q second=%q", firstKeys[len(firstKeys)-1], secondKeys[0]) + } + seen := make(map[string]struct{}, len(firstKeys)) + for _, key := range firstKeys { + seen[key] = struct{}{} + } + for _, key := range secondKeys { + if _, exists := seen[key]; exists { + b.Fatalf("cluster continuation repeated key %q", key) + } + } + } +} diff --git a/storage/benchio/read_gate.go b/storage/benchio/read_gate.go index 15f1d305..4f7795e7 100644 --- a/storage/benchio/read_gate.go +++ b/storage/benchio/read_gate.go @@ -18,6 +18,13 @@ func WaitForReadBudget(_ int) error { return nil } +// WaitForReadBudgetCancelable is the context-free build's no-op equivalent of +// the benchmark read gate. Production reads use their own cancellation +// boundary; this helper exists so benchmark wrappers share one interface. +func WaitForReadBudgetCancelable(_ <-chan struct{}, _ int) error { + return nil +} + // WrapPayloadReaderForBenchmark leaves readers unchanged outside benchmark builds. func WrapPayloadReaderForBenchmark(reader io.Reader) io.Reader { return reader @@ -27,3 +34,46 @@ func WrapPayloadReaderForBenchmark(reader io.Reader) io.Reader { func WrapPayloadReaderAtForBenchmark(reader io.ReaderAt) io.ReaderAt { return reader } + +// ResetPayloadStatsForBenchmark is a no-op outside benchmark builds. +func ResetPayloadStatsForBenchmark() {} + +// MarkPayloadCancellationForBenchmark is a no-op outside benchmark builds. +func MarkPayloadCancellationForBenchmark() {} + +// RecordListRowForBenchmark is a no-op outside benchmark builds. +func RecordListRowForBenchmark() {} + +// RecordPayloadBytesForBenchmark is a no-op outside benchmark builds. +func RecordPayloadBytesForBenchmark(_ int) {} + +// BeginListScanForBenchmark returns a no-op cleanup outside benchmark builds. +func BeginListScanForBenchmark() func() { + return func() {} +} + +// BeginListHandlerForBenchmark returns a no-op cleanup outside benchmark builds. +func BeginListHandlerForBenchmark() func() { + return func() {} +} + +// RecordPayloadReaderOpenedForBenchmark is a no-op outside benchmark builds. +func RecordPayloadReaderOpenedForBenchmark() {} + +// RecordPayloadReaderClosedForBenchmark is a no-op outside benchmark builds. +func RecordPayloadReaderClosedForBenchmark() {} + +// PostCancellationRowsForBenchmark returns zero outside benchmark builds. +func PostCancellationRowsForBenchmark() int64 { return 0 } + +// PostCancellationBytesForBenchmark returns zero outside benchmark builds. +func PostCancellationBytesForBenchmark() int64 { return 0 } + +// ActiveListScansForBenchmark returns zero outside benchmark builds. +func ActiveListScansForBenchmark() int64 { return 0 } + +// ActiveListHandlersForBenchmark returns zero outside benchmark builds. +func ActiveListHandlersForBenchmark() int64 { return 0 } + +// ActivePayloadReadersForBenchmark returns zero outside benchmark builds. +func ActivePayloadReadersForBenchmark() int64 { return 0 } diff --git a/storage/benchio/read_gate_benchmark.go b/storage/benchio/read_gate_benchmark.go index 6d692c6f..fe1fcd45 100644 --- a/storage/benchio/read_gate_benchmark.go +++ b/storage/benchio/read_gate_benchmark.go @@ -7,17 +7,108 @@ package benchio import ( "context" + "fmt" "io" "sync" + "sync/atomic" + "time" "golang.org/x/time/rate" ) const readGateBurst = 64 * 1024 +type readBlock struct { + started chan struct{} + release chan struct{} + startOnce sync.Once + releaseOnce sync.Once +} + var readGate struct { sync.RWMutex limiter *rate.Limiter + block *readBlock +} + +var benchmarkReadStats struct { + canceled atomic.Bool + postCancelRows atomic.Int64 + postCancelBytes atomic.Int64 + activeScans atomic.Int64 + activeHandlers atomic.Int64 + activeReaders atomic.Int64 +} + +// ResetPayloadStatsForBenchmark clears the counters for one measured list page. +func ResetPayloadStatsForBenchmark() { + benchmarkReadStats.canceled.Store(false) + benchmarkReadStats.postCancelRows.Store(0) + benchmarkReadStats.postCancelBytes.Store(0) + benchmarkReadStats.activeScans.Store(0) + benchmarkReadStats.activeHandlers.Store(0) + benchmarkReadStats.activeReaders.Store(0) +} + +// MarkPayloadCancellationForBenchmark marks the request cancellation boundary. +func MarkPayloadCancellationForBenchmark() { + benchmarkReadStats.canceled.Store(true) +} + +// RecordListRowForBenchmark records rows visited after cancellation. +func RecordListRowForBenchmark() { + if benchmarkReadStats.canceled.Load() { + benchmarkReadStats.postCancelRows.Add(1) + } +} + +// RecordPayloadBytesForBenchmark records payload bytes returned after cancellation. +func RecordPayloadBytesForBenchmark(bytes int) { + if benchmarkReadStats.canceled.Load() { + benchmarkReadStats.postCancelBytes.Add(int64(bytes)) + } +} + +// BeginListScanForBenchmark tracks the lifetime of a storage list iterator. +func BeginListScanForBenchmark() func() { + benchmarkReadStats.activeScans.Add(1) + return func() { benchmarkReadStats.activeScans.Add(-1) } +} + +// BeginListHandlerForBenchmark tracks one public or peer list handler. +func BeginListHandlerForBenchmark() func() { + benchmarkReadStats.activeHandlers.Add(1) + return func() { benchmarkReadStats.activeHandlers.Add(-1) } +} + +// RecordPayloadReaderOpenedForBenchmark tracks a foreground payload reader. +func RecordPayloadReaderOpenedForBenchmark() { + benchmarkReadStats.activeReaders.Add(1) +} + +// RecordPayloadReaderClosedForBenchmark tracks a closed foreground payload reader. +func RecordPayloadReaderClosedForBenchmark() { + benchmarkReadStats.activeReaders.Add(-1) +} + +func PostCancellationRowsForBenchmark() int64 { + return benchmarkReadStats.postCancelRows.Load() +} + +func PostCancellationBytesForBenchmark() int64 { + return benchmarkReadStats.postCancelBytes.Load() +} + +func ActiveListScansForBenchmark() int64 { + return benchmarkReadStats.activeScans.Load() +} + +func ActiveListHandlersForBenchmark() int64 { + return benchmarkReadStats.activeHandlers.Load() +} + +func ActivePayloadReadersForBenchmark() int64 { + return benchmarkReadStats.activeReaders.Load() } // SetReadRateLimitForBenchmark installs one shared payload-read budget for the @@ -39,12 +130,60 @@ func SetReadRateLimitForBenchmark(bytesPerSecond int64) func() { } } +// BlockPayloadReadsForBenchmark holds every foreground payload read at the +// read boundary until Release is called or the reader's interrupt channel is +// closed. It returns the first-read notification, a release function, and a +// restore function for the previous gate state. +func BlockPayloadReadsForBenchmark() (started <-chan struct{}, release func(), restore func()) { + block := &readBlock{ + started: make(chan struct{}), + release: make(chan struct{}), + } + + readGate.Lock() + previous := readGate.block + readGate.block = block + readGate.Unlock() + + release = func() { + block.releaseOnce.Do(func() { close(block.release) }) + } + restore = func() { + readGate.Lock() + if readGate.block == block { + readGate.block = previous + } + readGate.Unlock() + } + return block.started, release, restore +} + // WaitForReadBudget admits a payload read in bounded chunks so every raw-file // and segment reader in the benchmark shares the same throughput cap. func WaitForReadBudget(bytes int) error { + return WaitForReadBudgetCancelable(nil, bytes) +} + +// WaitForReadBudgetCancelable waits at the benchmark payload boundary while +// also allowing the private list reader to wake it through cancellation. +func WaitForReadBudgetCancelable(cancel <-chan struct{}, bytes int) error { readGate.RLock() + block := readGate.block limiter := readGate.limiter readGate.RUnlock() + + if block != nil { + block.startOnce.Do(func() { close(block.started) }) + if cancel == nil { + <-block.release + } else { + select { + case <-block.release: + case <-cancel: + return context.Canceled + } + } + } if limiter == nil { return nil } @@ -54,8 +193,27 @@ func WaitForReadBudget(bytes int) error { if chunk > readGateBurst { chunk = readGateBurst } - if err := limiter.WaitN(context.Background(), chunk); err != nil { - return err + + reservation := limiter.ReserveN(time.Now(), chunk) + if !reservation.OK() { + return fmt.Errorf("read budget reservation exceeds burst") + } + delay := reservation.Delay() + if delay > 0 { + timer := time.NewTimer(delay) + if cancel == nil { + <-timer.C + } else { + select { + case <-timer.C: + case <-cancel: + if !timer.Stop() { + <-timer.C + } + reservation.CancelAt(time.Now()) + return context.Canceled + } + } } bytes -= chunk } diff --git a/storage/fd/interruptible_reader.go b/storage/fd/interruptible_reader.go new file mode 100644 index 00000000..9125bed2 --- /dev/null +++ b/storage/fd/interruptible_reader.go @@ -0,0 +1,74 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package fd + +import ( + "io" + "os" + "sync" +) + +// InterruptibleReadCloser owns a private file descriptor. Interrupt closes only +// that descriptor; Close also runs the optional ownership callback. The +// descriptor must not be shared with another reader. Callers that need a cached +// descriptor should continue to use the regular FdCache APIs instead. +type InterruptibleReadCloser struct { + io.ReadSeeker + file *os.File + onClose func() + + interruptOnce sync.Once + closeOnce sync.Once +} + +// NewInterruptibleReadCloser creates a read closer for a private descriptor. +// onClose is called exactly once by Close, after the descriptor close is +// initiated. Callers should use Interrupt before Close when a read is blocked. +func NewInterruptibleReadCloser(file *os.File, reader io.ReadSeeker, onClose func()) *InterruptibleReadCloser { + return &InterruptibleReadCloser{ + ReadSeeker: reader, + file: file, + onClose: onClose, + } +} + +// Interrupt closes the private descriptor. Any read which outlives the close +// cannot affect another reader's cached descriptor because the descriptor is +// not shared. +func (r *InterruptibleReadCloser) Interrupt() { + if r == nil { + return + } + r.interruptOnce.Do(func() { + if r.file != nil { + _ = r.file.Close() + } + }) +} + +// ReadAll reads the private descriptor synchronously. List storage pairs this +// method with a context cancellation callback that interrupts the descriptor; +// callers using an arbitrary reader should use the asynchronous fallback +// instead. +func (r *InterruptibleReadCloser) ReadAll() ([]byte, error) { + return io.ReadAll(r) +} + +// Close is idempotent so it can safely race with Interrupt. It runs the +// optional ownership callback only once, even if a private read is still +// unwinding. +func (r *InterruptibleReadCloser) Close() error { + if r == nil { + return nil + } + r.closeOnce.Do(func() { + if r.file != nil { + _ = r.file.Close() + } + if r.onClose != nil { + r.onClose() + } + }) + return nil +} diff --git a/storage/files/manager.go b/storage/files/manager.go index 540dc3ac..80a5d605 100644 --- a/storage/files/manager.go +++ b/storage/files/manager.go @@ -151,6 +151,27 @@ func (fm *FileManager) Read(filePath string, length int64) (io.ReadCloser, error }), nil } +// ReadForList opens a private descriptor for a paginated value read. The +// file lock keeps deletion from winning between the metadata lookup and open; +// once the descriptor exists, cancellation can close it without touching the +// descriptor cache or another reader's lock. +func (fm *FileManager) ReadForList(filePath string, length int64) (io.ReadCloser, error) { + if filePath == "" || length <= 0 { + return nil, fmt.Errorf("invalid file path or length: path=%s, length=%d", filePath, length) + } + + fileLock := fd.GetFileLockManager().GetFileLock(filePath) + fileLock.RLock() + privateFile, err := os.Open(filePath) + fileLock.RUnlock() + if err != nil { + return nil, err + } + + reader := io.NewSectionReader(privateFile, 0, length) + return wrapReadForBenchmark(fd.NewInterruptibleReadCloser(privateFile, reader, nil)), nil +} + // Remove removes a file for the given key without blocking if it's being read func (fm *FileManager) Remove(filePath string) error { // Get file-specific lock (use GetFileLock, not Acquire, to avoid opening the file) diff --git a/storage/files/read_gate_benchmark.go b/storage/files/read_gate_benchmark.go index 9603471a..d3c6639e 100644 --- a/storage/files/read_gate_benchmark.go +++ b/storage/files/read_gate_benchmark.go @@ -8,20 +8,51 @@ package files import ( "fmt" "io" + "sync" "github.com/tigrisdata/ocache/storage/benchio" ) type benchmarkReadCloser struct { io.ReadCloser - seeker io.Seeker + seeker io.Seeker + interruptCh chan struct{} + interruptOnce sync.Once + closeOnce sync.Once } func (r *benchmarkReadCloser) Read(p []byte) (int, error) { - if err := benchio.WaitForReadBudget(len(p)); err != nil { + if err := benchio.WaitForReadBudgetCancelable(r.interruptCh, len(p)); err != nil { return 0, err } - return r.ReadCloser.Read(p) + n, err := r.ReadCloser.Read(p) + benchio.RecordPayloadBytesForBenchmark(n) + return n, err +} + +// Interrupt wakes benchmark-gated reads before forwarding the interrupt to +// the private production descriptor. +func (r *benchmarkReadCloser) Interrupt() { + r.interruptOnce.Do(func() { close(r.interruptCh) }) + if interrupter, ok := r.ReadCloser.(interface{ Interrupt() }); ok { + interrupter.Interrupt() + } +} + +// ReadAll preserves the synchronous private-file fast path while the storage +// list reader's cancellation callback interrupts this wrapper when needed. +func (r *benchmarkReadCloser) ReadAll() ([]byte, error) { + return io.ReadAll(r) +} + +func (r *benchmarkReadCloser) Close() error { + var err error + r.closeOnce.Do(func() { + benchio.RecordPayloadReaderClosedForBenchmark() + r.Interrupt() + err = r.ReadCloser.Close() + }) + return err } func (r *benchmarkReadCloser) Seek(offset int64, whence int) (int64, error) { @@ -32,6 +63,11 @@ func (r *benchmarkReadCloser) Seek(offset int64, whence int) (int64, error) { } func wrapReadForBenchmark(reader io.ReadCloser) io.ReadCloser { + benchio.RecordPayloadReaderOpenedForBenchmark() readerWithSeek, _ := reader.(io.Seeker) - return &benchmarkReadCloser{ReadCloser: reader, seeker: readerWithSeek} + return &benchmarkReadCloser{ + ReadCloser: reader, + seeker: readerWithSeek, + interruptCh: make(chan struct{}), + } } diff --git a/storage/list_values_cancellation_test.go b/storage/list_values_cancellation_test.go new file mode 100644 index 00000000..5dda3937 --- /dev/null +++ b/storage/list_values_cancellation_test.go @@ -0,0 +1,102 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package storage + +import ( + "context" + "errors" + "io" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type blockingListReader struct { + started chan struct{} + finish chan struct{} + closed chan struct{} + done chan struct{} + finishOnce sync.Once + closeOnce sync.Once +} + +func newBlockingListReader() *blockingListReader { + return &blockingListReader{ + started: make(chan struct{}), + finish: make(chan struct{}), + closed: make(chan struct{}), + done: make(chan struct{}), + } +} + +func (r *blockingListReader) Read([]byte) (int, error) { + select { + case <-r.started: + default: + close(r.started) + } + <-r.finish + close(r.done) + return 0, errors.New("injected read completion") +} + +func (r *blockingListReader) Close() error { + r.closeOnce.Do(func() { close(r.closed) }) + return nil +} + +func (r *blockingListReader) Interrupt() {} + +func (r *blockingListReader) finishRead() { + r.finishOnce.Do(func() { close(r.finish) }) +} + +func TestReadListPayloadCancelsAndClosesNeverReturningRead(t *testing.T) { + reader := newBlockingListReader() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + payloadReader := newListPayloadReader(ctx) + defer payloadReader.close() + + resultCh := make(chan error, 1) + go func() { + _, err := payloadReader.read(ctx, reader) + resultCh <- err + }() + + select { + case <-reader.started: + case <-time.After(time.Second): + t.Fatal("injected payload read did not start") + } + cancel() + + select { + case err := <-resultCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("payload read did not return after cancellation") + } + + select { + case <-reader.closed: + case <-time.After(time.Second): + t.Fatal("payload reader was not closed after cancellation") + } + + // The injected read ignores Interrupt and Close until its source is + // released. The list operation must still return after closing its private + // reader, while the worker is allowed to finish independently. + reader.finishRead() + select { + case <-reader.done: + case <-time.After(time.Second): + t.Fatal("injected payload reader did not finish cleanup") + } +} + +var _ io.ReadCloser = (*blockingListReader)(nil) diff --git a/storage/list_values_cap_test.go b/storage/list_values_cap_test.go index ccfb49d7..082fc69f 100644 --- a/storage/list_values_cap_test.go +++ b/storage/list_values_cap_test.go @@ -5,6 +5,7 @@ package storage import ( "bytes" + "context" "testing" "github.com/stretchr/testify/assert" @@ -30,7 +31,7 @@ func TestListKeyValuesWithPagination_OmitsOversizeValues(t *testing.T) { require.NoError(t, s.Put("a-small", bytes.NewReader(small), 0)) require.NoError(t, s.Put("b-large", bytes.NewReader(large), 0)) - entries, _, _, err := s.ListKeyValuesWithPagination("", "", 100) + entries, _, _, err := s.ListKeyValuesWithPagination(context.Background(), "", "", 100) require.NoError(t, err) got := make(map[string]KeyValue, len(entries)) @@ -48,6 +49,64 @@ func TestListKeyValuesWithPagination_OmitsOversizeValues(t *testing.T) { assert.Equal(t, int64(len(large)), got["b-large"].ValueLength, "size must still be reported") } +// TestListKeyValuesWithPaginationReadsRawValuesAndPreservesCursor verifies +// that spilled raw-file values retain page ordering and continuation behavior. +func TestListKeyValuesWithPaginationReadsRawValuesAndPreservesCursor(t *testing.T) { + s, cleanup := createTestStorage(t, 0, 1, 4*1024, 16*1024*1024, 1000, 0) + defer cleanup() + + value := bytes.Repeat([]byte("r"), 8*1024) + for _, key := range []string{"list-01", "list-02", "list-03"} { + require.NoError(t, s.Put(key, bytes.NewReader(value), 0)) + } + + entries, lastKey, hasMore, err := s.ListKeyValuesWithPagination(context.Background(), "list-", "", 2) + require.NoError(t, err) + require.Equal(t, []string{"list-01", "list-02"}, []string{entries[0].Key, entries[1].Key}) + require.Equal(t, "list-02", lastKey) + require.True(t, hasMore) + for _, entry := range entries { + require.Equal(t, value, entry.Value) + require.Equal(t, int64(len(value)), entry.ValueLength) + require.False(t, entry.ValueOmitted) + } + + entries, lastKey, hasMore, err = s.ListKeyValuesWithPagination(context.Background(), "list-", lastKey, 2) + require.NoError(t, err) + require.Equal(t, []string{"list-03"}, []string{entries[0].Key}) + require.Empty(t, lastKey) + require.False(t, hasMore) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + entries, lastKey, hasMore, err = s.ListKeyValuesWithPagination(ctx, "list-", "", 2) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, entries) + require.Empty(t, lastKey) + require.False(t, hasMore) +} + +// TestListKeyValuesWithPaginationReadsSegmentValues verifies that the +// interruptible list reader also returns values from compacted segments. +func TestListKeyValuesWithPaginationReadsSegmentValues(t *testing.T) { + s, cleanup := createTestStorage(t, 0, 1, 64*1024, 16*1024*1024, 1000, 0) + defer cleanup() + + value := bytes.Repeat([]byte("s"), 8*1024) + require.NoError(t, s.Put("segment-list", bytes.NewReader(value), 0)) + _, _ = s.compactor.CompactFiles(context.Background(), 0) + + entries, lastKey, hasMore, err := s.ListKeyValuesWithPagination(context.Background(), "segment-", "", 1) + require.NoError(t, err) + require.Empty(t, lastKey) + require.False(t, hasMore) + require.Len(t, entries, 1) + require.Equal(t, "segment-list", entries[0].Key) + require.Equal(t, value, entries[0].Value) + require.Equal(t, int64(len(value)), entries[0].ValueLength) + require.False(t, entries[0].ValueOmitted) +} + // TestNewStorageWithConfig_ClampsCompactThresholdBelowSegmentSize verifies the // startup guardrail clamps an invalid CompactThreshold (>= SegmentSize) rather // than failing to start (#165). diff --git a/storage/segment/manager.go b/storage/segment/manager.go index 18f63afa..93cde0c7 100644 --- a/storage/segment/manager.go +++ b/storage/segment/manager.go @@ -102,6 +102,36 @@ func (sm *Manager) ReadEntry(userKey string, segPath string, offset, length int6 return wrapReadForBenchmark(reader), nil } +// ReadEntryForList opens a private descriptor for a paginated value read. The +// segment and file locks cover the lookup-to-open window, then cancellation can +// close the private descriptor without touching the shared descriptor cache. +func (sm *Manager) ReadEntryForList(userKey string, segPath string, offset, length int64) (io.ReadCloser, error) { + if segPath == "" || offset < 0 || length <= 0 { + return nil, fmt.Errorf("invalid segment path, offset or length: path=%s, offset=%d, length=%d", segPath, offset, length) + } + + sm.mu.RLock() + seg := sm.segMap[segPath] + sm.mu.RUnlock() + if seg == nil { + return nil, fmt.Errorf("segment not found: %s", segPath) + } + + fileLock := fd.GetFileLockManager().GetFileLock(seg.path) + fileLock.RLock() + seg.mu.RLock() + privateFile, err := os.Open(seg.path) + seg.mu.RUnlock() + fileLock.RUnlock() + if err != nil { + return nil, err + } + + offset += CalculateValueHeaderSize(userKey) + reader := io.NewSectionReader(privateFile, offset, length) + return wrapReadForBenchmark(fd.NewInterruptibleReadCloser(privateFile, reader, nil)), nil +} + // AcquireOpenSegmentWithReservation returns an open segment reserved for the caller // The callerID should be unique per goroutine/thread (e.g., "compactor", "recompactor-1") // The segment will be reserved exclusively for this caller until released diff --git a/storage/segment/read_gate_benchmark.go b/storage/segment/read_gate_benchmark.go index 169175d3..399d1037 100644 --- a/storage/segment/read_gate_benchmark.go +++ b/storage/segment/read_gate_benchmark.go @@ -8,20 +8,51 @@ package segment import ( "fmt" "io" + "sync" "github.com/tigrisdata/ocache/storage/benchio" ) type benchmarkReadCloser struct { io.ReadCloser - seeker io.Seeker + seeker io.Seeker + interruptCh chan struct{} + interruptOnce sync.Once + closeOnce sync.Once } func (r *benchmarkReadCloser) Read(p []byte) (int, error) { - if err := benchio.WaitForReadBudget(len(p)); err != nil { + if err := benchio.WaitForReadBudgetCancelable(r.interruptCh, len(p)); err != nil { return 0, err } - return r.ReadCloser.Read(p) + n, err := r.ReadCloser.Read(p) + benchio.RecordPayloadBytesForBenchmark(n) + return n, err +} + +// Interrupt wakes benchmark-gated reads before forwarding the interrupt to +// the private production descriptor. +func (r *benchmarkReadCloser) Interrupt() { + r.interruptOnce.Do(func() { close(r.interruptCh) }) + if interrupter, ok := r.ReadCloser.(interface{ Interrupt() }); ok { + interrupter.Interrupt() + } +} + +// ReadAll preserves the synchronous private-file fast path while the storage +// list reader's cancellation callback interrupts this wrapper when needed. +func (r *benchmarkReadCloser) ReadAll() ([]byte, error) { + return io.ReadAll(r) +} + +func (r *benchmarkReadCloser) Close() error { + var err error + r.closeOnce.Do(func() { + benchio.RecordPayloadReaderClosedForBenchmark() + r.Interrupt() + err = r.ReadCloser.Close() + }) + return err } func (r *benchmarkReadCloser) Seek(offset int64, whence int) (int64, error) { @@ -32,6 +63,11 @@ func (r *benchmarkReadCloser) Seek(offset int64, whence int) (int64, error) { } func wrapReadForBenchmark(reader io.ReadCloser) io.ReadCloser { + benchio.RecordPayloadReaderOpenedForBenchmark() readerWithSeek, _ := reader.(io.Seeker) - return &benchmarkReadCloser{ReadCloser: reader, seeker: readerWithSeek} + return &benchmarkReadCloser{ + ReadCloser: reader, + seeker: readerWithSeek, + interruptCh: make(chan struct{}), + } } diff --git a/storage/storage.go b/storage/storage.go index eee53e15..644cbe49 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -5,6 +5,7 @@ package storage import ( "bytes" + "context" "errors" "fmt" "io" @@ -20,6 +21,7 @@ import ( "github.com/tigrisdata/ocache/common/bufferpool" "github.com/tigrisdata/ocache/common/metrics" + "github.com/tigrisdata/ocache/storage/benchio" "github.com/tigrisdata/ocache/storage/compaction" "github.com/tigrisdata/ocache/storage/deletion" storageErrors "github.com/tigrisdata/ocache/storage/errors" @@ -617,16 +619,213 @@ type KeyValue struct { ValueOmitted bool } +// interruptibleListReader is implemented by the private readers used for +// paginated values. Interrupt stops the private read when the reader supports +// it; Close also releases the shared file ownership without closing a cached +// descriptor. +type interruptibleListReader interface { + io.ReadCloser + Interrupt() +} + +type listPayloadReadResult struct { + data []byte + err error +} + +type listPayloadReadAller interface { + ReadAll() ([]byte, error) +} + +// listPayloadReader uses a synchronous fast path for the private file readers +// owned by this package. A context callback interrupts the current descriptor, +// so the handler does not need a per-value worker or channel round trip. An +// arbitrary interruptible reader uses the worker fallback because its Read may +// ignore Interrupt and Close; that path can outlive the canceled list without +// retaining shared descriptor ownership. +type listPayloadReader struct { + requests chan io.ReadCloser + results chan listPayloadReadResult + stop chan struct{} + stopOnce sync.Once + workerOnce sync.Once + + currentMu sync.Mutex + current interruptibleListReader + stopWatch func() bool +} + +func newListPayloadReader(ctx context.Context) *listPayloadReader { + reader := &listPayloadReader{stop: make(chan struct{})} + reader.stopWatch = context.AfterFunc(ctx, reader.interruptCurrent) + return reader +} + +func (r *listPayloadReader) interruptCurrent() { + r.currentMu.Lock() + defer r.currentMu.Unlock() + if r.current != nil { + r.current.Interrupt() + _ = r.current.Close() + } +} + +func (r *listPayloadReader) setCurrent(ctx context.Context, reader interruptibleListReader) error { + r.currentMu.Lock() + defer r.currentMu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + r.current = reader + return nil +} + +func (r *listPayloadReader) clearCurrent() { + r.currentMu.Lock() + r.current = nil + r.currentMu.Unlock() +} + +func (r *listPayloadReader) startWorker() { + r.workerOnce.Do(func() { + r.requests = make(chan io.ReadCloser, 1) + r.results = make(chan listPayloadReadResult, 1) + go r.run() + }) +} + +func (r *listPayloadReader) run() { + for { + select { + case reader := <-r.requests: + data, err := io.ReadAll(reader) + _ = reader.Close() + select { + case r.results <- listPayloadReadResult{data: data, err: err}: + case <-r.stop: + return + } + case <-r.stop: + return + } + } +} + +func (r *listPayloadReader) close() { + if r.stopWatch != nil { + r.stopWatch() + } + r.stopOnce.Do(func() { close(r.stop) }) +} + +func (r *listPayloadReader) readSynchronous(ctx context.Context, reader io.ReadCloser, readAll func() ([]byte, error)) ([]byte, error) { + interruptible, ok := reader.(interruptibleListReader) + if !ok { + _ = reader.Close() + return nil, fmt.Errorf("list payload reader is not interruptible") + } + if err := r.setCurrent(ctx, interruptible); err != nil { + _ = reader.Close() + return nil, err + } + + data, err := readAll() + r.clearCurrent() + _ = reader.Close() + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return data, err +} + +func (r *listPayloadReader) read(ctx context.Context, reader io.ReadCloser) ([]byte, error) { + interruptible, ok := reader.(interruptibleListReader) + if !ok { + _ = reader.Close() + return nil, fmt.Errorf("list payload reader is not interruptible") + } + if err := ctx.Err(); err != nil { + _ = reader.Close() + return nil, err + } + r.startWorker() + + select { + case r.requests <- reader: + case <-ctx.Done(): + interruptible.Interrupt() + _ = reader.Close() + return nil, ctx.Err() + } + + select { + case result := <-r.results: + _ = reader.Close() + if err := ctx.Err(); err != nil { + return nil, err + } + return result.data, result.err + case <-ctx.Done(): + // The reader owns only a private descriptor. Interrupt and Close can + // therefore release the private descriptor immediately; a blocked read + // cannot invalidate another reader's cached descriptor. + interruptible.Interrupt() + _ = reader.Close() + return nil, ctx.Err() + } +} + +func (s *Storage) readListPayload(ctx context.Context, payloadReader *listPayloadReader, valueKey string, valueMsg *pb.ValueMessage) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + var ( + reader io.ReadCloser + readErr error + ) + switch valueMsg.ValueType { + case pb.ValueType_SEGMENT: + reader, readErr = s.segmentManager.ReadEntryForList(valueKey, valueMsg.SegmentPath, valueMsg.SegmentOffset, valueMsg.ValueLength) + case pb.ValueType_RAW_FILE: + reader, readErr = s.fileManager.ReadForList(valueMsg.RawFilePath, valueMsg.ValueLength) + default: + return nil, nil + } + if readErr != nil || reader == nil { + return nil, readErr + } + if ctx.Done() == nil { + data, err := io.ReadAll(reader) + _ = reader.Close() + return data, err + } + if readAller, ok := reader.(listPayloadReadAller); ok { + return payloadReader.readSynchronous(ctx, reader, readAller.ReadAll) + } + return payloadReader.read(ctx, reader) +} + // ListKeyValuesWithPagination returns paginated, sorted key-value pairs from RocksDB. // For inline values the data is returned directly; for file/segment values the data is // read from disk. Returns: (entries, lastKey, hasMore, error). -func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string, limit int) ([]KeyValue, string, bool, error) { +func (s *Storage) ListKeyValuesWithPagination(ctx context.Context, userPrefix string, startKey string, limit int) ([]KeyValue, string, bool, error) { + if ctx == nil { + ctx = context.Background() + } + endBenchmarkScan := benchio.BeginListScanForBenchmark() + defer endBenchmarkScan() + storageType := "unknown" start := time.Now() defer func() { metrics.StorageOperationDuration.WithLabelValues("list_kv_paginated", storageType).Observe(float64(time.Since(start).Milliseconds())) }() + if err := ctx.Err(); err != nil { + return nil, "", false, err + } + ro := metadata.CreateReadOptions(true, false) defer ro.Destroy() it := s.meta.Handle().NewIterator(ro) @@ -634,6 +833,12 @@ func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string var entries []KeyValue var lastKey string + var payloadReader *listPayloadReader + defer func() { + if payloadReader != nil { + payloadReader.close() + } + }() // Determine where to start iteration var seekKey []byte @@ -655,26 +860,52 @@ func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string // Seek to start position it.Seek(seekKey) + if err := ctx.Err(); err != nil { + return nil, "", false, err + } // If we have a startKey, skip it (pagination is exclusive) if startKey != "" && it.Valid() { - k := it.Key().Data() - currentUserKey := keys.ExtractUserKey(k) + keySlice := it.Key() + valueSlice := it.Value() + currentUserKey := keys.ExtractUserKey(keySlice.Data()) + keySlice.Free() + valueSlice.Free() if currentUserKey == startKey { - it.Key().Free() - it.Value().Free() + if err := ctx.Err(); err != nil { + return nil, "", false, err + } it.Next() } } - // Collect up to limit key-value pairs - for it.ValidForPrefix(prefixBoundary) { + // Collect up to limit key-value pairs. A cancellation discards the entire + // page so callers never receive a cursor that describes a partial scan. + for { + if err := ctx.Err(); err != nil { + return nil, "", false, err + } + if !it.ValidForPrefix(prefixBoundary) { + break + } if limit > 0 && len(entries) >= limit { break } + benchio.RecordListRowForBenchmark() + + keySlice := it.Key() + valueSlice := it.Value() + k := keySlice.Data() + v := valueSlice.Data() + releaseSlices := func() { + keySlice.Free() + valueSlice.Free() + } - k := it.Key().Data() - v := it.Value().Data() + if err := ctx.Err(); err != nil { + releaseSlices() + return nil, "", false, err + } // Try to decode as proto ValueMessage to check expiry. // If unmarshal fails, include the key with nil value to keep @@ -685,8 +916,10 @@ func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string valueMsg := &pb.ValueMessage{} if err := proto.Unmarshal(v, valueMsg); err == nil { if valueMsg.Expiry > 0 && time.Now().Unix() >= valueMsg.Expiry { - it.Key().Free() - it.Value().Free() + releaseSlices() + if err := ctx.Err(); err != nil { + return nil, "", false, err + } it.Next() continue } @@ -709,41 +942,40 @@ func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string case pb.ValueType_INLINE: // proto.Unmarshal owns a copy separate from the iterator buffer. data = valueMsg.Data - case pb.ValueType_SEGMENT: - r, readErr := s.segmentManager.ReadEntry(keys.ExtractUserKey(k), valueMsg.SegmentPath, valueMsg.SegmentOffset, valueMsg.ValueLength) - if readErr == nil && r != nil { - data, readErr = io.ReadAll(r) - if closer, ok := r.(io.Closer); ok { - closer.Close() - } - if readErr != nil { - data = nil - } + case pb.ValueType_SEGMENT, pb.ValueType_RAW_FILE: + if payloadReader == nil && ctx.Done() != nil { + payloadReader = newListPayloadReader(ctx) } - case pb.ValueType_RAW_FILE: - r, readErr := s.fileManager.Read(valueMsg.RawFilePath, valueMsg.ValueLength) - if readErr == nil && r != nil { - data, readErr = io.ReadAll(r) - if closer, ok := r.(io.Closer); ok { - closer.Close() - } - if readErr != nil { - data = nil + var readErr error + data, readErr = s.readListPayload(ctx, payloadReader, keys.ExtractUserKey(k), valueMsg) + if readErr != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + releaseSlices() + return nil, "", false, ctxErr } + // Preserve the existing best-effort behavior for ordinary + // payload read failures: keep the key and report no value. + data = nil } } } } + if err := ctx.Err(); err != nil { + releaseSlices() + return nil, "", false, err + } userKey := keys.ExtractUserKey(k) entries = append(entries, KeyValue{Key: userKey, Value: data, ValueLength: valueLength, ValueOmitted: valueOmitted}) lastKey = userKey - it.Key().Free() - it.Value().Free() + releaseSlices() it.Next() } + if err := ctx.Err(); err != nil { + return nil, "", false, err + } if err := it.Err(); err != nil { metrics.StorageOperations.WithLabelValues("list_kv_paginated", storageType, "error").Inc() metrics.Errors.WithLabelValues("rocksdb", "list_kv_paginated").Inc() @@ -753,12 +985,15 @@ func (s *Storage) ListKeyValuesWithPagination(userPrefix string, startKey string // Check if there are more keys hasMore := it.ValidForPrefix(prefixBoundary) if hasMore { - k := it.Key().Data() - nextUserKey := keys.ExtractUserKey(k) + keySlice := it.Key() + nextUserKey := keys.ExtractUserKey(keySlice.Data()) if userPrefix != "" && !bytes.HasPrefix([]byte(nextUserKey), []byte(userPrefix)) { hasMore = false } - it.Key().Free() + keySlice.Free() + } + if err := ctx.Err(); err != nil { + return nil, "", false, err } if !hasMore { diff --git a/tests/integration/test_utils.go b/tests/integration/test_utils.go index ab2f5350..35588cfd 100644 --- a/tests/integration/test_utils.go +++ b/tests/integration/test_utils.go @@ -5,6 +5,7 @@ package integration import ( "bytes" + "context" "fmt" "io" "math" @@ -256,7 +257,7 @@ func (h *IntegrationTestHarness) ListPage(prefix string, limit int, continuation // ListPageWithValues returns a page of key-value pairs with pagination support func (h *IntegrationTestHarness) ListPageWithValues(prefix string, limit int, continuationToken string) (entries []KeyValueEntry, nextToken string, hasMore bool, err error) { - storageEntries, lastKey, more, listErr := h.Storage.ListKeyValuesWithPagination(prefix, continuationToken, limit) + storageEntries, lastKey, more, listErr := h.Storage.ListKeyValuesWithPagination(context.Background(), prefix, continuationToken, limit) if listErr != nil { h.Metrics.ErrorCount.Add(1) return nil, "", false, listErr