Skip to content

Commit 17f2d53

Browse files
trunk-io[bot]trunk-io[bot]
authored
ptstorage: lock the meta row during protect and release (#170655)
ptstorage: lock the meta row during protect and release
2 parents ff68e02 + eae678c commit 17f2d53

3 files changed

Lines changed: 210 additions & 6 deletions

File tree

‎pkg/kv/kvserver/protectedts/ptstorage/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ go_test(
3434
name = "ptstorage_test",
3535
size = "large",
3636
srcs = [
37+
"bench_test.go",
3738
"main_test.go",
3839
"metrics_test.go",
3940
"storage_test.go",
@@ -63,6 +64,7 @@ go_test(
6364
"//pkg/testutils/serverutils",
6465
"//pkg/testutils/skip",
6566
"//pkg/testutils/testcluster",
67+
"//pkg/util/ctxgroup",
6668
"//pkg/util/hlc",
6769
"//pkg/util/leaktest",
6870
"//pkg/util/log",

‎pkg/kv/kvserver/protectedts/ptstorage/bench_test.go‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
// Copyright 2026 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package ptstorage_test
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"sync/atomic"
12+
"testing"
13+
14+
"github.com/cockroachdb/cockroach/pkg/base"
15+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
16+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
17+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptstorage"
18+
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
19+
"github.com/cockroachdb/cockroach/pkg/sql/isql"
20+
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
21+
"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
22+
"github.com/cockroachdb/cockroach/pkg/util/hlc"
23+
"github.com/cockroachdb/cockroach/pkg/util/log"
24+
"github.com/cockroachdb/cockroach/pkg/util/uuid"
25+
"github.com/stretchr/testify/require"
26+
)
27+
28+
// workerCounts is the parallelism sweep shared by the throughput benchmarks
29+
// in this file.
30+
// var workerCounts = []int{1, 2, 4, 8, 16, 32, 64, 128}
31+
var workerCounts = []int{128}
32+
33+
// startBenchServer brings up a single-node in-process server suitable for the
34+
// throughput benchmarks.
35+
func startBenchServer(b *testing.B) (protectedts.Storage, func()) {
36+
ctx := context.Background()
37+
srv := serverutils.StartServerOnly(b, base.TestServerArgs{})
38+
s := srv.ApplicationLayer()
39+
40+
protectedts.MaxBytes.Override(ctx, &s.ClusterSettings().SV, 0)
41+
protectedts.MaxSpans.Override(ctx, &s.ClusterSettings().SV, 0)
42+
43+
pts := ptstorage.WithDatabase(
44+
ptstorage.New(s.ClusterSettings(), &protectedts.TestingKnobs{}),
45+
s.InternalDB().(isql.DB),
46+
)
47+
return pts, func() { srv.Stopper().Stop(ctx) }
48+
}
49+
50+
// runConcurrent invokes op b.N times spread across `workers` goroutines and
51+
// then waits for them all to finish. op receives the per-invocation index in
52+
// [0, b.N), which is useful for indexing into a per-iteration data set.
53+
//
54+
// Errors returned by op are counted and reported as the "errs/op" metric
55+
// rather than failing the benchmark. Under heavy contention some operations
56+
// can exhaust the internal kv retry budget; counting rather than failing lets
57+
// us still produce a comparable baseline number.
58+
func runConcurrent(b *testing.B, ctx context.Context, workers int, op func(idx int) error) {
59+
var (
60+
next atomic.Int64
61+
errs atomic.Int64
62+
)
63+
64+
b.ResetTimer()
65+
err := ctxgroup.GroupWorkers(ctx, workers, func(context.Context, int) error {
66+
for {
67+
idx := int(next.Add(1) - 1)
68+
if idx >= b.N {
69+
return nil
70+
}
71+
if err := op(idx); err != nil {
72+
errs.Add(1)
73+
}
74+
}
75+
})
76+
require.NoError(b, err)
77+
b.StopTimer()
78+
b.ReportMetric(float64(errs.Load())/float64(b.N), "errs/op")
79+
}
80+
81+
// BenchmarkProtect measures the throughput of protectedts.Storage.Protect
82+
// calls as a function of writer concurrency. Each iteration writes one new
83+
// record. Per-op wall-clock time reported by the framework is the inverse of
84+
// sustained throughput: ops/sec = 1e9 / ns/op.
85+
func BenchmarkProtect(b *testing.B) {
86+
defer log.Scope(b).Close(b)
87+
ctx := context.Background()
88+
89+
pts, stop := startBenchServer(b)
90+
defer stop()
91+
92+
now := hlc.Timestamp{WallTime: 1}
93+
target := ptpb.MakeSchemaObjectsTarget([]descpb.ID{42})
94+
95+
for _, workers := range workerCounts {
96+
b.Run(fmt.Sprintf("workers=%d", workers), func(b *testing.B) {
97+
runConcurrent(b, ctx, workers, func(_ int) error {
98+
rec := ptpb.Record{
99+
ID: uuid.MakeV4().GetBytes(),
100+
Timestamp: now,
101+
Mode: ptpb.PROTECT_AFTER,
102+
Target: target,
103+
}
104+
return pts.Protect(ctx, &rec)
105+
})
106+
})
107+
}
108+
}
109+
110+
// BenchmarkRelease measures the throughput of protectedts.Storage.Release
111+
// calls as a function of writer concurrency. Each sub-benchmark first
112+
// pre-populates b.N records (the setup time is excluded from the
113+
// measurement) and then releases each one exactly once.
114+
//
115+
// Like Protect, Release is a read-modify-write against the singleton meta
116+
// row, so it shares Protect's contention shape on that row.
117+
func BenchmarkRelease(b *testing.B) {
118+
defer log.Scope(b).Close(b)
119+
ctx := context.Background()
120+
121+
pts, stop := startBenchServer(b)
122+
defer stop()
123+
124+
now := hlc.Timestamp{WallTime: 1}
125+
target := ptpb.MakeSchemaObjectsTarget([]descpb.ID{42})
126+
127+
for _, workers := range workerCounts {
128+
b.Run(fmt.Sprintf("workers=%d", workers), func(b *testing.B) {
129+
ids := prePopulate(b, ctx, pts, workers, b.N, now, target)
130+
runConcurrent(b, ctx, workers, func(idx int) error {
131+
return pts.Release(ctx, ids[idx])
132+
})
133+
})
134+
}
135+
}
136+
137+
// prePopulate inserts n protected timestamp records in parallel using
138+
// `workers` goroutines and returns their IDs. It is intended for use as
139+
// benchmark setup (before b.ResetTimer); a Protect failure here is fatal
140+
// because it means the benchmark has nothing meaningful to measure.
141+
func prePopulate(
142+
b *testing.B,
143+
ctx context.Context,
144+
pts protectedts.Storage,
145+
workers, n int,
146+
ts hlc.Timestamp,
147+
target *ptpb.Target,
148+
) []uuid.UUID {
149+
ids := make([]uuid.UUID, n)
150+
var next atomic.Int64
151+
err := ctxgroup.GroupWorkers(ctx, workers, func(context.Context, int) error {
152+
for {
153+
idx := int(next.Add(1) - 1)
154+
if idx >= n {
155+
return nil
156+
}
157+
id := uuid.MakeV4()
158+
rec := ptpb.Record{
159+
ID: id.GetBytes(),
160+
Timestamp: ts,
161+
Mode: ptpb.PROTECT_AFTER,
162+
Target: target,
163+
}
164+
if err := pts.Protect(ctx, &rec); err != nil {
165+
return err
166+
}
167+
ids[idx] = id
168+
}
169+
})
170+
require.NoError(b, err)
171+
return ids
172+
}

‎pkg/kv/kvserver/protectedts/ptstorage/sql.go‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,14 @@ package ptstorage
77

88
const (
99

10-
// currentMetaCTE is used by all queries which access the meta row.
11-
// The query returns a default row if there currently is no meta row.
12-
// At the time of writing, there will never be a physical row in the meta
13-
// table with version zero.
10+
// currentMetaCTE reads the meta row, returning a default zero row if
11+
// none exists yet. At the time of writing, there will never be a
12+
// physical row in the meta table with version zero.
13+
//
14+
// This is the read-only variant, used by getMetadataQuery. Mutating
15+
// queries (protect, release) use currentMetaCTEForUpdate to avoid the
16+
// WriteTooOld retry storm that concurrent writers would otherwise
17+
// trigger on the shared meta row.
1418
currentMetaCTE = `
1519
SELECT
1620
version, num_records, num_spans, total_bytes
@@ -24,9 +28,35 @@ LIMIT
2428
1
2529
`
2630

31+
// currentMetaCTEForUpdate is the locking variant of currentMetaCTE used by
32+
// the protect and release queries. Both are read-modify-writes against
33+
// the singleton meta row: every concurrent caller reads the same row,
34+
// then upserts a new version. Without serialization the writers race,
35+
// producing a storm of WriteTooOld retries on the shared key. Acquiring
36+
// an exclusive lock on the real meta row during the read forces those
37+
// callers to queue rather than collide. The FOR UPDATE clause is bound
38+
// to the real-table leg only; the synthetic zero-row fallback (used when
39+
// no meta row exists yet) does not lock anything.
40+
currentMetaCTEForUpdate = `
41+
SELECT
42+
version, num_records, num_spans, total_bytes
43+
FROM
44+
(
45+
SELECT version, num_records, num_spans, total_bytes
46+
FROM system.protected_ts_meta
47+
FOR UPDATE
48+
)
49+
UNION ALL
50+
SELECT 0 AS version, 0 AS num_records, 0 AS num_spans, 0 AS total_bytes
51+
ORDER BY
52+
version DESC
53+
LIMIT
54+
1
55+
`
56+
2757
protectQuery = `
2858
WITH
29-
current_meta AS (` + currentMetaCTE + `),
59+
current_meta AS (` + currentMetaCTEForUpdate + `),
3060
checks AS (` + protectChecksCTE + `),
3161
updated_meta AS (` + protectUpsertMetaCTE + `),
3262
new_record AS (` + protectInsertRecordCTEWithChecks + `)
@@ -114,7 +144,7 @@ RETURNING
114144

115145
releaseQueryWithMeta = `
116146
WITH
117-
current_meta AS (` + currentMetaCTE + `),
147+
current_meta AS (` + currentMetaCTEForUpdate + `),
118148
record AS (` + releaseSelectRecordCTE + `),
119149
updated_meta AS (` + releaseUpsertMetaCTE + `)
120150
DELETE FROM

0 commit comments

Comments
 (0)