From 0c5d3dd587a614b0c681030b4cec372a21dc6463 Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Fri, 17 Jul 2026 17:15:52 -0400 Subject: [PATCH 1/4] Add Bigtable historical state offload: client and Kafka ingestion (part 1) Part 1 of the Bigtable historical MVCC state offload: the Bigtable data client (hand-rolled over the repo's pinned gRPC v1.57, with bounded read retries, keepalive, and billing-dimension cost metrics), the frozen row key encoding (m|shard|store|key|inverted-version mutation rows and v|bucket|inverted-version version markers), the point-read Reader, and the Kafka consumer that ingests the changelog stream into Bigtable with per-partition ordering, version dedupe, locality-chunked bulk writes, and offset commits only after durable sink writes. Connects to Google Cloud Managed Kafka via TLS + SASL/PLAIN. Part 2 wires the node-side pruned-read fallback on top of this. Co-Authored-By: Claude Fable 5 --- go.mod | 3 + go.sum | 6 + sei-db/state_db/ss/offload/consumer/README.md | 57 ++ .../state_db/ss/offload/consumer/bigtable.go | 233 ++++++ .../ss/offload/consumer/bigtable_test.go | 266 +++++++ .../cmd/historical-offload-consumer/main.go | 68 ++ .../state_db/ss/offload/consumer/compact.go | 63 ++ .../ss/offload/consumer/compact_test.go | 65 ++ sei-db/state_db/ss/offload/consumer/config.go | 76 ++ .../consumer/config/example-bigtable.json | 25 + .../ss/offload/consumer/config_test.go | 54 ++ .../state_db/ss/offload/consumer/consumer.go | 287 +++++++ sei-db/state_db/ss/offload/consumer/kafka.go | 126 +++ .../ss/offload/consumer/kafka_test.go | 77 ++ .../state_db/ss/offload/consumer/metrics.go | 66 ++ .../ss/offload/consumer/metrics_test.go | 36 + sei-db/state_db/ss/offload/consumer/sink.go | 21 + .../ss/offload/historical/bigtable.go | 749 ++++++++++++++++++ .../ss/offload/historical/bigtable_test.go | 242 ++++++ .../state_db/ss/offload/historical/metrics.go | 135 ++++ .../ss/offload/historical/metrics_test.go | 51 ++ .../state_db/ss/offload/historical/reader.go | 28 + sei-db/state_db/ss/offload/kafka.go | 57 +- 23 files changed, 2772 insertions(+), 19 deletions(-) create mode 100644 sei-db/state_db/ss/offload/consumer/README.md create mode 100644 sei-db/state_db/ss/offload/consumer/bigtable.go create mode 100644 sei-db/state_db/ss/offload/consumer/bigtable_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go create mode 100644 sei-db/state_db/ss/offload/consumer/compact.go create mode 100644 sei-db/state_db/ss/offload/consumer/compact_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/config.go create mode 100644 sei-db/state_db/ss/offload/consumer/config/example-bigtable.json create mode 100644 sei-db/state_db/ss/offload/consumer/config_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/consumer.go create mode 100644 sei-db/state_db/ss/offload/consumer/kafka.go create mode 100644 sei-db/state_db/ss/offload/consumer/kafka_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/metrics.go create mode 100644 sei-db/state_db/ss/offload/consumer/metrics_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/sink.go create mode 100644 sei-db/state_db/ss/offload/historical/bigtable.go create mode 100644 sei-db/state_db/ss/offload/historical/bigtable_test.go create mode 100644 sei-db/state_db/ss/offload/historical/metrics.go create mode 100644 sei-db/state_db/ss/offload/historical/metrics_test.go create mode 100644 sei-db/state_db/ss/offload/historical/reader.go diff --git a/go.mod b/go.mod index 66bfbda919..3136e045f3 100644 --- a/go.mod +++ b/go.mod @@ -137,6 +137,8 @@ require ( ) require ( + cloud.google.com/go/bigtable v1.37.0 + cloud.google.com/go/compute/metadata v0.8.0 // indirect filippo.io/edwards25519 v1.1.0 // indirect github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 // indirect github.com/DataDog/zstd v1.5.7 // indirect @@ -265,6 +267,7 @@ require ( go.yaml.in/yaml/v3 v3.0.4 // indirect golang.org/x/image v0.36.0 golang.org/x/mod v0.32.0 + golang.org/x/oauth2 v0.30.0 // indirect golang.org/x/term v0.39.0 // indirect golang.org/x/text v0.34.0 // indirect google.golang.org/genproto v0.0.0-20250603155806-513f23925822 // indirect diff --git a/go.sum b/go.sum index 86a857402a..3aa7bf9eb9 100644 --- a/go.sum +++ b/go.sum @@ -125,6 +125,8 @@ cloud.google.com/go/bigquery v1.47.0/go.mod h1:sA9XOgy0A8vQK9+MWhEQTY6Tix87M/Zur cloud.google.com/go/bigquery v1.48.0/go.mod h1:QAwSz+ipNgfL5jxiaK7weyOhzdoAy1zFm0Nf1fysJac= cloud.google.com/go/bigquery v1.49.0/go.mod h1:Sv8hMmTFFYBlt/ftw2uN6dFdQPzBlREY9yBh7Oy7/4Q= cloud.google.com/go/bigquery v1.50.0/go.mod h1:YrleYEh2pSEbgTBZYMJ5SuSr0ML3ypjRB1zgf7pvQLU= +cloud.google.com/go/bigtable v1.37.0 h1:Q+x7y04lQ0B+WXp03wc1/FLhFt4CwcQdkwWT0M4Jp3w= +cloud.google.com/go/bigtable v1.37.0/go.mod h1:HXqddP6hduwzrtiTCqZPpj9ij4hGZb4Zy1WF/dT+yaU= cloud.google.com/go/billing v1.4.0/go.mod h1:g9IdKBEFlItS8bTtlrZdVLWSSdSyFUZKXNS02zKMOZY= cloud.google.com/go/billing v1.5.0/go.mod h1:mztb1tBc3QekhjSgmpf/CV4LzWXLzCArwpLmP2Gm88s= cloud.google.com/go/billing v1.6.0/go.mod h1:WoXzguj+BeHXPbKfNWkqVtDdzORazmCjraY+vrxcyvI= @@ -175,6 +177,8 @@ cloud.google.com/go/compute/metadata v0.1.0/go.mod h1:Z1VN+bulIf6bt4P/C37K4DyZYZ cloud.google.com/go/compute/metadata v0.2.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k= cloud.google.com/go/compute/metadata v0.2.1/go.mod h1:jgHgmJd2RKBGzXqF5LR2EZMGxBkeanZ9wwa75XHJgOM= cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA= +cloud.google.com/go/compute/metadata v0.8.0 h1:HxMRIbao8w17ZX6wBnjhcDkW6lTFpgcaobyVfZWqRLA= +cloud.google.com/go/compute/metadata v0.8.0/go.mod h1:sYOGTp851OV9bOFJ9CH7elVvyzopvWQFNNghtDQ/Biw= cloud.google.com/go/contactcenterinsights v1.3.0/go.mod h1:Eu2oemoePuEFc/xKFPjbTuPSj0fYJcPls9TFlPNnHHY= cloud.google.com/go/contactcenterinsights v1.4.0/go.mod h1:L2YzkGbPsv+vMQMCADxJoT9YiTTnSEd6fEvCeHTYVck= cloud.google.com/go/contactcenterinsights v1.6.0/go.mod h1:IIDlT6CLcDoyv79kDv8iWxMSTZhLxSCofVV5W6YFM/w= @@ -2287,6 +2291,8 @@ golang.org/x/oauth2 v0.0.0-20221014153046-6fdb5e3db783/go.mod h1:h4gKUeWbJ4rQPri golang.org/x/oauth2 v0.5.0/go.mod h1:9/XBHVqLaWO3/BRHs5jbpYCnOZVjj5V0ndyaAM7KB4I= golang.org/x/oauth2 v0.6.0/go.mod h1:ycmewcwgD4Rpr3eZJLSB4Kyyljb3qDh40vJ8STE5HKw= golang.org/x/oauth2 v0.7.0/go.mod h1:hPLQkd9LyjfXTiRohC/41GhcFqxisoUQ99sCUOHO9x4= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= golang.org/x/perf v0.0.0-20230113213139-801c7ef9e5c5/go.mod h1:UBKtEnL8aqnd+0JHqZ+2qoMDwtuy6cYhhKNoHLBiTQc= golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= diff --git a/sei-db/state_db/ss/offload/consumer/README.md b/sei-db/state_db/ss/offload/consumer/README.md new file mode 100644 index 0000000000..7e2ec5cef1 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/README.md @@ -0,0 +1,57 @@ +# Historical State Offload (Bigtable) + +Bigtable holds immutable MVCC mutation rows for history that local SS has +pruned. The shape is narrow: + +- local SS remains the hot store for recent state, writes, imports, pruning, and iterators +- Bigtable keeps immutable MVCC mutation rows for older history +- reads below local SS retention can fall back to Bigtable for `Get` and `Has` + +Row keys are salted with an inverted height suffix: + +```text +m | shard(store,key) | store_name | state_key | inverted_height +``` + +Reads scan from `inverted(target_height)` and stop after the first row, giving +the latest write at or before the requested height. Ordered prefix iteration is +intentionally not served from the offload store. + +## Consumer + +The consumer reads historical offload changelog messages from Kafka and writes +them into Bigtable. Kafka offsets are committed only after the sink write +succeeds. Mutation rows are written before the version marker. + +```bash +cbt -project my-gcp-project -instance sei-history createtable state_mutations +cbt -project my-gcp-project -instance sei-history createfamily state_mutations state + +go run ./sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer \ + ./sei-db/state_db/ss/offload/consumer/config/example-bigtable.json +``` + +The example config is a local/dev placeholder. Set real Kafka brokers and +Bigtable credentials/config in your own config. + +For Google Cloud Managed Service for Apache Kafka, connect with TLS plus +SASL/PLAIN using service-account credentials: + +```json +"Kafka": { + "Brokers": ["bootstrap.CLUSTER.REGION.managedkafka.PROJECT.cloud.goog:9092"], + "TLSEnabled": true, + "SASLMechanism": "plain", + "Username": "kafka-client@PROJECT.iam.gserviceaccount.com", + "Password": "" +} +``` + +## Current Limits + +- The node-side read fallback lands in part 2; this part is the client library + and the ingestion pipeline. +- No cross-row transaction on ingest; mutation rows are written first and the + version marker is written last, so replay is idempotent after partial failure. +- No automatic table creation from the binary. +- No backfill tooling; coverage starts when ingestion starts. diff --git a/sei-db/state_db/ss/offload/consumer/bigtable.go b/sei-db/state_db/ss/offload/consumer/bigtable.go new file mode 100644 index 0000000000..0e70ce46a3 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/bigtable.go @@ -0,0 +1,233 @@ +package consumer + +import ( + "context" + "fmt" + "sort" + "strconv" + "time" + + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" + "golang.org/x/sync/errgroup" +) + +type BigtableConfig = historical.BigtableConfig + +const ( + defaultBigtableMutationChunkRows = 1024 + defaultBigtableMutationChunkConcurrency = 8 +) + +type bigtableSink struct { + client *historical.BigtableClient + applyBulk historical.BigtableApplyBulkFunc + family string + shards int + bulkChunkRows int + bulkChunkWorkers int +} + +var _ Sink = (*bigtableSink)(nil) + +func NewBigtableSink(cfg BigtableConfig) (Sink, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + ctx := context.Background() + client, err := historical.OpenBigtableClient(ctx, cfg) + if err != nil { + return nil, err + } + return &bigtableSink{ + client: client, + applyBulk: client.ApplyBulk, + family: cfg.Family, + shards: cfg.Shards, + bulkChunkRows: defaultBigtableMutationChunkRows, + bulkChunkWorkers: defaultBigtableMutationChunkConcurrency, + }, nil +} + +func (s *bigtableSink) Close() error { + if s.client != nil { + return s.client.Close() + } + return nil +} + +func (s *bigtableSink) WriteBatch(ctx context.Context, records []Record) error { + records = compactRecords(records) + if len(records) == 0 { + return nil + } + if err := s.writeRecordRows(ctx, records); err != nil { + return err + } + return s.writeVersionMarkers(ctx, records) +} + +func (s *bigtableSink) writeRecordRows(ctx context.Context, records []Record) error { + rows := make([]historical.BigtableRowMutation, 0, bigtableRowMutationCount(records)) + for _, rec := range records { + rows = s.appendRecordRowMutations(rows, rec.Entry) + } + if len(rows) == 0 { + return nil + } + return s.applyRecordRowMutations(ctx, rows) +} + +func (s *bigtableSink) applyRecordRowMutations(ctx context.Context, rows []historical.BigtableRowMutation) error { + chunks := bigtableRowMutationChunks(rows, s.bulkChunkRows) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(s.bulkChunkWorkers) + for _, chunk := range chunks { + chunk := chunk + g.Go(func() error { + errs, err := s.applyBulk(gctx, chunk) + return bigtableBulkError(chunk, errs, err) + }) + } + return g.Wait() +} + +func (s *bigtableSink) appendRecordRowMutations(rows []historical.BigtableRowMutation, entry *proto.ChangelogEntry) []historical.BigtableRowMutation { + for _, mutation := range compactMutations(entry) { + rows = append(rows, s.mutationRow(entry.Version, mutation.storeName, mutation.pair)) + } + for _, up := range entry.Upgrades { + rows = append(rows, s.upgradeRow(entry.Version, up)) + } + return rows +} + +// mutationRow writes value+deleted cells for live pairs but only a deleted +// cell for tombstones, saving a cell per delete. Readers must therefore check +// the deleted column before trusting any value cell — a replayed live write +// followed by a tombstone leaves both cells on the row. +func (s *bigtableSink) mutationRow(version int64, storeName string, pair *proto.KVPair) historical.BigtableRowMutation { + ts := historical.BigtableTimestamp(version) + deleted := pair.Delete || pair.Value == nil + rowKey := historical.BigtableMutationRowKey(storeName, pair.Key, version, s.shards) + if deleted { + return historical.BigtableRowMutation{ + RowKey: rowKey, + SetCells: []historical.BigtableSetCell{{ + Family: s.family, + Qualifier: historical.BigtableDeletedColumn, + TimestampMicros: ts, + Value: boolByte(true), + }}, + } + } + return historical.BigtableRowMutation{ + RowKey: rowKey, + SetCells: []historical.BigtableSetCell{ + {Family: s.family, Qualifier: historical.BigtableValueColumn, TimestampMicros: ts, Value: pair.Value}, + {Family: s.family, Qualifier: historical.BigtableDeletedColumn, TimestampMicros: ts, Value: boolByte(false)}, + }, + } +} + +func (s *bigtableSink) upgradeRow(version int64, up *proto.TreeNameUpgrade) historical.BigtableRowMutation { + ts := historical.BigtableTimestamp(version) + return historical.BigtableRowMutation{ + RowKey: historical.BigtableUpgradeRowKey(version, up.Name), + SetCells: []historical.BigtableSetCell{ + {Family: s.family, Qualifier: "rename_from", TimestampMicros: ts, Value: []byte(up.RenameFrom)}, + {Family: s.family, Qualifier: historical.BigtableDeletedColumn, TimestampMicros: ts, Value: boolByte(up.Delete)}, + }, + } +} + +func (s *bigtableSink) writeVersionMarkers(ctx context.Context, records []Record) error { + rows := make([]historical.BigtableRowMutation, 0, len(records)) + ingestedAt := []byte(strconv.FormatInt(time.Now().UnixNano(), 10)) + for _, rec := range records { + version := rec.Entry.Version + ts := historical.BigtableTimestamp(version) + rows = append(rows, historical.BigtableRowMutation{ + RowKey: historical.BigtableVersionRowKey(version), + SetCells: []historical.BigtableSetCell{ + {Family: s.family, Qualifier: "topic", TimestampMicros: ts, Value: []byte(rec.Topic)}, + {Family: s.family, Qualifier: "partition", TimestampMicros: ts, Value: []byte(strconv.Itoa(rec.Partition))}, + {Family: s.family, Qualifier: "offset", TimestampMicros: ts, Value: []byte(strconv.FormatInt(rec.Offset, 10))}, + {Family: s.family, Qualifier: "ingested_at_unix_nano", TimestampMicros: ts, Value: ingestedAt}, + }, + }) + } + errs, err := s.applyBulk(ctx, rows) + if err := bigtableBulkError(rows, errs, err); err != nil { + return fmt.Errorf("insert bigtable version markers: %w", err) + } + return nil +} + +func bigtableRowMutationCount(records []Record) int { + total := 0 + for _, rec := range records { + total += entryMutationCapacity(rec.Entry) + len(rec.Entry.Upgrades) + } + return total +} + +func bigtableRowMutationChunks(rows []historical.BigtableRowMutation, maxRows int) [][]historical.BigtableRowMutation { + if len(rows) == 0 { + return nil + } + if maxRows <= 0 { + maxRows = len(rows) + } + sort.Slice(rows, func(i, j int) bool { + return rows[i].RowKey < rows[j].RowKey + }) + + chunks := make([][]historical.BigtableRowMutation, 0, (len(rows)+maxRows-1)/maxRows) + start := 0 + startLocality := bigtableRowLocality(rows[0].RowKey) + for i := 1; i < len(rows); i++ { + locality := bigtableRowLocality(rows[i].RowKey) + if i-start >= maxRows || locality != startLocality { + chunks = append(chunks, rows[start:i]) + start = i + startLocality = locality + } + } + return append(chunks, rows[start:]) +} + +func bigtableRowLocality(rowKey string) string { + // Mutation row keys are m|shard|store|key|version; keep chunks inside one + // shard prefix so separate chunks can hit separate Bigtable tablets. + if len(rowKey) >= 3 && rowKey[0] == 'm' { + return rowKey[:3] + } + if len(rowKey) > 0 { + return rowKey[:1] + } + return rowKey +} + +func bigtableBulkError(rows []historical.BigtableRowMutation, errs []error, err error) error { + if err != nil { + return err + } + if len(errs) != len(rows) { + return fmt.Errorf("bigtable returned %d mutation results for %d rows", len(errs), len(rows)) + } + for i, rowErr := range errs { + if rowErr != nil { + return fmt.Errorf("row %q: %w", rows[i].RowKey, rowErr) + } + } + return nil +} + +func boolByte(v bool) []byte { + if v { + return []byte{1} + } + return []byte{0} +} diff --git a/sei-db/state_db/ss/offload/consumer/bigtable_test.go b/sei-db/state_db/ss/offload/consumer/bigtable_test.go new file mode 100644 index 0000000000..63bae0d7b2 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/bigtable_test.go @@ -0,0 +1,266 @@ +package consumer + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" + "github.com/stretchr/testify/require" +) + +func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { + var rows []string + sink := &bigtableSink{ + family: historical.DefaultBigtableFamily, + shards: historical.DefaultBigtableShards, + bulkChunkWorkers: 1, + applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + for _, mutation := range mutations { + rows = append(rows, mutation.RowKey) + } + return make([]error, len(mutations)), nil + }, + } + entry := &proto.ChangelogEntry{ + Version: 7, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{ + {Key: []byte("k1"), Value: []byte("old")}, + {Key: []byte("k1"), Value: []byte("new")}, + {Key: []byte("drop"), Delete: true}, + }}, + }}, + Upgrades: []*proto.TreeNameUpgrade{{Name: "new-store"}}, + } + + require.NoError(t, sink.WriteBatch(context.Background(), []Record{{Topic: "t", Partition: 1, Offset: 2, Entry: entry}})) + require.Len(t, rows, 4) + require.ElementsMatch(t, []string{ + historical.BigtableMutationRowKey("bank", []byte("k1"), 7, historical.DefaultBigtableShards), + historical.BigtableMutationRowKey("bank", []byte("drop"), 7, historical.DefaultBigtableShards), + historical.BigtableUpgradeRowKey(7, "new-store"), + }, rows[:3]) + require.Equal(t, historical.BigtableVersionRowKey(7), rows[3]) +} + +func TestBigtableSinkWriteBatchWritesRowsBeforeMarkers(t *testing.T) { + var calls [][]string + + sink := &bigtableSink{ + family: historical.DefaultBigtableFamily, + shards: historical.DefaultBigtableShards, + bulkChunkWorkers: 1, + applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + call := make([]string, 0, len(mutations)) + for _, mutation := range mutations { + call = append(call, mutation.RowKey) + } + calls = append(calls, call) + return make([]error, len(mutations)), nil + }, + } + records := []Record{ + {Entry: &proto.ChangelogEntry{ + Version: 1, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{{Key: []byte("k1"), Value: []byte("v1")}}}, + }}, + }}, + {Entry: &proto.ChangelogEntry{ + Version: 2, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{{Key: []byte("k2"), Value: []byte("v2")}}}, + }}, + }}, + } + + require.NoError(t, sink.WriteBatch(context.Background(), records)) + require.GreaterOrEqual(t, len(calls), 2) + require.ElementsMatch(t, []string{ + historical.BigtableMutationRowKey("bank", []byte("k1"), 1, historical.DefaultBigtableShards), + historical.BigtableMutationRowKey("bank", []byte("k2"), 2, historical.DefaultBigtableShards), + }, flattenCalls(calls[:len(calls)-1])) + require.Equal(t, []string{ + historical.BigtableVersionRowKey(1), + historical.BigtableVersionRowKey(2), + }, calls[len(calls)-1]) +} + +func TestBigtableRowMutationChunksSortsAndGroupsByLocality(t *testing.T) { + rows := []historical.BigtableRowMutation{ + {RowKey: string([]byte{'u', 'z'})}, + {RowKey: string([]byte{'m', 0, 2, 'b'})}, + {RowKey: string([]byte{'m', 0, 1, 'c'})}, + {RowKey: string([]byte{'m', 0, 1, 'a'})}, + {RowKey: string([]byte{'m', 0, 1, 'b'})}, + {RowKey: string([]byte{'m', 0, 2, 'a'})}, + } + + chunks := bigtableRowMutationChunks(rows, 2) + + require.Equal(t, [][]string{ + { + string([]byte{'m', 0, 1, 'a'}), + string([]byte{'m', 0, 1, 'b'}), + }, + { + string([]byte{'m', 0, 1, 'c'}), + }, + { + string([]byte{'m', 0, 2, 'a'}), + string([]byte{'m', 0, 2, 'b'}), + }, + { + string([]byte{'u', 'z'}), + }, + }, chunkRowKeys(chunks)) +} + +func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { + var calls [][]string + sink := &bigtableSink{ + family: historical.DefaultBigtableFamily, + shards: 2, + bulkChunkRows: 1, + bulkChunkWorkers: 1, + applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + call := make([]string, 0, len(mutations)) + for _, mutation := range mutations { + call = append(call, mutation.RowKey) + } + calls = append(calls, call) + return make([]error, len(mutations)), nil + }, + } + records := []Record{ + {Entry: &proto.ChangelogEntry{ + Version: 1, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{ + {Key: []byte("k1"), Value: []byte("v1")}, + {Key: []byte("k2"), Value: []byte("v2")}, + }}, + }}, + Upgrades: []*proto.TreeNameUpgrade{{Name: "new-store"}}, + }}, + {Entry: &proto.ChangelogEntry{ + Version: 2, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{ + {Key: []byte("k3"), Value: []byte("v3")}, + }}, + }}, + }}, + } + + require.NoError(t, sink.WriteBatch(context.Background(), records)) + + require.Greater(t, len(calls), 2) + markerCall := calls[len(calls)-1] + require.Equal(t, []string{ + historical.BigtableVersionRowKey(1), + historical.BigtableVersionRowKey(2), + }, markerCall) + for _, call := range calls[:len(calls)-1] { + require.Len(t, call, 1) + require.NotEqual(t, markerCall, call) + } + require.ElementsMatch(t, []string{ + historical.BigtableMutationRowKey("bank", []byte("k1"), 1, 2), + historical.BigtableMutationRowKey("bank", []byte("k2"), 1, 2), + historical.BigtableUpgradeRowKey(1, "new-store"), + historical.BigtableMutationRowKey("bank", []byte("k3"), 2, 2), + }, flattenCalls(calls[:len(calls)-1])) +} + +func TestBigtableSinkAppliesRecordChunksConcurrently(t *testing.T) { + started := make(chan struct{}, 2) + release := make(chan struct{}) + sink := &bigtableSink{ + bulkChunkRows: 1, + bulkChunkWorkers: 2, + applyBulk: func(ctx context.Context, rows []historical.BigtableRowMutation) ([]error, error) { + select { + case started <- struct{}{}: + case <-ctx.Done(): + return nil, ctx.Err() + } + select { + case <-release: + return make([]error, len(rows)), nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + rows := []historical.BigtableRowMutation{ + {RowKey: string([]byte{'m', 0, 1, 'a'})}, + {RowKey: string([]byte{'m', 0, 2, 'a'})}, + } + + errCh := make(chan error, 1) + go func() { + errCh <- sink.applyRecordRowMutations(context.Background(), rows) + }() + + for i := 0; i < 2; i++ { + select { + case <-started: + case <-time.After(time.Second): + close(release) + require.FailNow(t, "timed out waiting for concurrent bigtable chunks") + } + } + close(release) + require.NoError(t, <-errCh) +} + +func TestBigtableBulkErrorValidatesMutationResultCount(t *testing.T) { + rows := []historical.BigtableRowMutation{{RowKey: "row-1"}, {RowKey: "row-2"}} + + err := bigtableBulkError(rows, []error{nil}, nil) + + require.ErrorContains(t, err, "mutation results") +} + +func TestBigtableBulkErrorWrapsRowError(t *testing.T) { + rowErr := errors.New("failed") + rows := []historical.BigtableRowMutation{{RowKey: "row-1"}} + + err := bigtableBulkError(rows, []error{rowErr}, nil) + + require.ErrorIs(t, err, rowErr) + require.ErrorContains(t, err, "row-1") +} + +func chunkRowKeys(chunks [][]historical.BigtableRowMutation) [][]string { + out := make([][]string, 0, len(chunks)) + for _, chunk := range chunks { + out = append(out, flattenRowMutations(chunk)) + } + return out +} + +func flattenCalls(calls [][]string) []string { + var out []string + for _, call := range calls { + out = append(out, call...) + } + return out +} + +func flattenRowMutations(rows []historical.BigtableRowMutation) []string { + out := make([]string, 0, len(rows)) + for _, row := range rows { + out = append(out, row.RowKey) + } + return out +} diff --git a/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go b/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go new file mode 100644 index 0000000000..1fc43636ce --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go @@ -0,0 +1,68 @@ +package main + +import ( + "context" + "fmt" + "log" + "os" + "os/signal" + "syscall" + "time" + + commonmetrics "github.com/sei-protocol/sei-chain/sei-db/common/metrics" + "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/consumer" +) + +func main() { + if len(os.Args) != 2 { + fmt.Fprintf(os.Stderr, "usage: %s \n", os.Args[0]) + os.Exit(2) + } + + cfg, err := consumer.LoadConfig(os.Args[1]) + if err != nil { + log.Fatalf("load config: %v", err) + } + + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer cancel() + + // Install the Prometheus MeterProvider before opening the sink so the + // backend cost metrics it registers bind to a real exporter. + if cfg.MetricsAddr != "" { + reg, shutdown, err := commonmetrics.SetupOtelPrometheus() + if err != nil { + log.Fatalf("setup metrics: %v", err) + } + defer func() { + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer shutdownCancel() + _ = shutdown(shutdownCtx) + }() + commonmetrics.StartMetricsServer(ctx, reg, cfg.MetricsAddr) + log.Printf("serving consumer metrics at %s/metrics", cfg.MetricsAddr) + } + + sink, err := consumer.NewBigtableSink(cfg.Bigtable) + if err != nil { + log.Fatalf("open bigtable sink: %v", err) + } + defer func() { _ = sink.Close() }() + + reader, err := consumer.NewKafkaReader(cfg.Kafka) + if err != nil { + log.Fatalf("open kafka reader: %v", err) + } + defer func() { _ = reader.Close() }() + + c := consumer.New(reader, sink, consumer.Options{ + Logf: func(format string, args ...interface{}) { log.Printf(format, args...) }, + Workers: cfg.Workers, + ShardBufferSize: cfg.ShardBufferSize, + MaxBatchRecords: cfg.MaxBatchRecords, + BatchMaxWait: time.Duration(cfg.BatchMaxWaitMS) * time.Millisecond, + }) + if err := c.Run(ctx); err != nil { + log.Fatalf("consumer: %v", err) + } +} diff --git a/sei-db/state_db/ss/offload/consumer/compact.go b/sei-db/state_db/ss/offload/consumer/compact.go new file mode 100644 index 0000000000..9f0495472e --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/compact.go @@ -0,0 +1,63 @@ +package consumer + +import "github.com/sei-protocol/sei-chain/sei-db/proto" + +type stateMutation struct { + storeName string + pair *proto.KVPair +} + +type stateMutationKey struct { + storeName string + key string +} + +// compactRecords drops records without entries and collapses records that +// share a changelog version (Kafka at-least-once redelivery) onto the latest +// one, so sinks write each version's rows once and version markers carry the +// newest offset. Order is preserved by the first occurrence of each version. +func compactRecords(records []Record) []Record { + indexByVersion := make(map[int64]int, len(records)) + out := make([]Record, 0, len(records)) + for _, rec := range records { + if rec.Entry == nil { + continue + } + if idx, ok := indexByVersion[rec.Entry.Version]; ok { + out[idx] = rec + continue + } + indexByVersion[rec.Entry.Version] = len(out) + out = append(out, rec) + } + return out +} + +func compactMutations(entry *proto.ChangelogEntry) []stateMutation { + if entry == nil { + return nil + } + mutations := make([]stateMutation, 0, entryMutationCapacity(entry)) + indexByKey := make(map[stateMutationKey]int, cap(mutations)) + for _, ncs := range entry.Changesets { + storeName := ncs.Name + for _, pair := range ncs.Changeset.Pairs { + key := stateMutationKey{storeName: storeName, key: string(pair.Key)} + if idx, ok := indexByKey[key]; ok { + mutations[idx].pair = pair + continue + } + indexByKey[key] = len(mutations) + mutations = append(mutations, stateMutation{storeName: storeName, pair: pair}) + } + } + return mutations +} + +func entryMutationCapacity(entry *proto.ChangelogEntry) int { + total := 0 + for _, changeset := range entry.Changesets { + total += len(changeset.Changeset.Pairs) + } + return total +} diff --git a/sei-db/state_db/ss/offload/consumer/compact_test.go b/sei-db/state_db/ss/offload/consumer/compact_test.go new file mode 100644 index 0000000000..919de6f9e0 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/compact_test.go @@ -0,0 +1,65 @@ +package consumer + +import ( + "testing" + + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/stretchr/testify/require" +) + +func TestCompactRecordsDropsNilEntries(t *testing.T) { + records := compactRecords([]Record{ + {Entry: &proto.ChangelogEntry{Version: 1}}, + {}, + {Entry: &proto.ChangelogEntry{Version: 2}}, + }) + require.Len(t, records, 2) + require.Equal(t, int64(1), records[0].Entry.Version) + require.Equal(t, int64(2), records[1].Entry.Version) +} + +func TestCompactRecordsCollapsesRedeliveredVersions(t *testing.T) { + records := compactRecords([]Record{ + {Offset: 10, Entry: &proto.ChangelogEntry{Version: 1}}, + {Offset: 11, Entry: &proto.ChangelogEntry{Version: 2}}, + {Offset: 12, Entry: &proto.ChangelogEntry{Version: 1}}, + {Offset: 13, Entry: &proto.ChangelogEntry{Version: 3}}, + }) + require.Len(t, records, 3) + // Version order is preserved; the redelivered version keeps its slot but + // carries the newest offset for the version marker. + require.Equal(t, int64(1), records[0].Entry.Version) + require.Equal(t, int64(12), records[0].Offset) + require.Equal(t, int64(2), records[1].Entry.Version) + require.Equal(t, int64(3), records[2].Entry.Version) +} + +func TestCompactMutationsKeepsLastWritePerKey(t *testing.T) { + entry := &proto.ChangelogEntry{ + Version: 9, + Changesets: []*proto.NamedChangeSet{{ + Name: "bank", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{ + {Key: []byte("k"), Value: []byte("old")}, + {Key: []byte("drop"), Value: []byte("present")}, + {Key: []byte("k"), Value: []byte("new")}, + {Key: []byte("drop"), Delete: true}, + }}, + }, { + Name: "evm", + Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{ + {Key: []byte("k"), Value: []byte("separate-store")}, + }}, + }}, + } + + mutations := compactMutations(entry) + require.Len(t, mutations, 3) + byKey := map[string]*proto.KVPair{} + for _, m := range mutations { + byKey[m.storeName+"/"+string(m.pair.Key)] = m.pair + } + require.Equal(t, []byte("new"), byKey["bank/k"].Value) + require.True(t, byKey["bank/drop"].Delete) + require.Equal(t, []byte("separate-store"), byKey["evm/k"].Value) +} diff --git a/sei-db/state_db/ss/offload/consumer/config.go b/sei-db/state_db/ss/offload/consumer/config.go new file mode 100644 index 0000000000..18889f6bf1 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/config.go @@ -0,0 +1,76 @@ +package consumer + +import ( + "encoding/json" + "fmt" + "os" +) + +const ( + defaultMaxBatchRecords = 128 + defaultBatchMaxWaitMS = 25 +) + +type Config struct { + Kafka KafkaReaderConfig + Bigtable BigtableConfig + Workers int + ShardBufferSize int + MaxBatchRecords int + BatchMaxWaitMS int + // MetricsAddr, when set (e.g. ":2112"), serves Prometheus metrics at + // /metrics so the backend cost counters (bigtable_rows_mutated_total, + // bigtable_bytes_written_total, bigtable_mutate_latency_seconds, ...) can be + // scraped. Empty disables the endpoint. + MetricsAddr string +} + +func (c *Config) Validate() error { + if err := c.Kafka.Validate(); err != nil { + return fmt.Errorf("kafka: %w", err) + } + bigtable := c.Bigtable + bigtable.ApplyDefaults() + if err := bigtable.Validate(); err != nil { + return fmt.Errorf("bigtable: %w", err) + } + if c.Workers < 0 { + return fmt.Errorf("workers must be non-negative") + } + if c.ShardBufferSize < 0 { + return fmt.Errorf("shard buffer size must be non-negative") + } + if c.MaxBatchRecords < 0 { + return fmt.Errorf("max batch records must be non-negative") + } + if c.BatchMaxWaitMS < 0 { + return fmt.Errorf("batch max wait ms must be non-negative") + } + return nil +} + +func (c *Config) applyDefaults() { + if c.MaxBatchRecords == 0 { + c.MaxBatchRecords = defaultMaxBatchRecords + } + if c.BatchMaxWaitMS == 0 { + c.BatchMaxWaitMS = defaultBatchMaxWaitMS + } +} + +func LoadConfig(path string) (*Config, error) { + // #nosec G304 -- config path is supplied by the operator on the command line. + raw, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read config: %w", err) + } + cfg := &Config{} + if err := json.Unmarshal(raw, cfg); err != nil { + return nil, fmt.Errorf("parse config: %w", err) + } + cfg.applyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/sei-db/state_db/ss/offload/consumer/config/example-bigtable.json b/sei-db/state_db/ss/offload/consumer/config/example-bigtable.json new file mode 100644 index 0000000000..47ba1f0d33 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/config/example-bigtable.json @@ -0,0 +1,25 @@ +{ + "Kafka": { + "Brokers": ["bootstrap.my-cluster.us-east1.managedkafka.my-gcp-project.cloud.goog:9092"], + "Topic": "historical-offload", + "GroupID": "historical-bigtable", + "StartOffset": "first", + "TLSEnabled": true, + "SASLMechanism": "plain", + "Username": "kafka-client@my-gcp-project.iam.gserviceaccount.com", + "Password": "" + }, + "Bigtable": { + "ProjectID": "my-gcp-project", + "InstanceID": "sei-history", + "Table": "state_mutations", + "Family": "state", + "AppProfile": "", + "Shards": 256 + }, + "Workers": 16, + "ShardBufferSize": 128, + "MaxBatchRecords": 128, + "BatchMaxWaitMS": 25, + "MetricsAddr": ":2112" +} diff --git a/sei-db/state_db/ss/offload/consumer/config_test.go b/sei-db/state_db/ss/offload/consumer/config_test.go new file mode 100644 index 0000000000..9cd6f65d3f --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/config_test.go @@ -0,0 +1,54 @@ +package consumer + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestConfigValidate(t *testing.T) { + cfg := Config{ + Kafka: KafkaReaderConfig{ + Brokers: []string{"localhost:9092"}, + Topic: "historical-offload", + GroupID: "historical-bigtable", + }, + Bigtable: BigtableConfig{ + ProjectID: "project", + InstanceID: "instance", + Table: "state", + }, + } + require.NoError(t, cfg.Validate()) + + cfg.Bigtable.Table = "" + require.ErrorContains(t, cfg.Validate(), "bigtable") +} + +func TestConfigApplyDefaults(t *testing.T) { + cfg := Config{} + cfg.applyDefaults() + require.Equal(t, defaultMaxBatchRecords, cfg.MaxBatchRecords) + require.Equal(t, defaultBatchMaxWaitMS, cfg.BatchMaxWaitMS) + + cfg = Config{MaxBatchRecords: 32, BatchMaxWaitMS: 5} + cfg.applyDefaults() + require.Equal(t, 32, cfg.MaxBatchRecords) + require.Equal(t, 5, cfg.BatchMaxWaitMS) +} + +func TestLoadConfig(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{ + "Kafka": {"Brokers": ["localhost:9092"], "Topic": "historical-offload", "GroupID": "g"}, + "Bigtable": {"ProjectID": "p", "InstanceID": "i", "Table": "t"}, + "MetricsAddr": ":2112" + }`), 0o600)) + + cfg, err := LoadConfig(path) + require.NoError(t, err) + require.Equal(t, ":2112", cfg.MetricsAddr) + require.Equal(t, defaultMaxBatchRecords, cfg.MaxBatchRecords) +} diff --git a/sei-db/state_db/ss/offload/consumer/consumer.go b/sei-db/state_db/ss/offload/consumer/consumer.go new file mode 100644 index 0000000000..d021cd9c6b --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/consumer.go @@ -0,0 +1,287 @@ +package consumer + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/segmentio/kafka-go" + "golang.org/x/sync/errgroup" +) + +// MessageSource is the subset of *kafka.Reader used by the loop. +type MessageSource interface { + FetchMessage(ctx context.Context) (kafka.Message, error) + CommitMessages(ctx context.Context, msgs ...kafka.Message) error +} + +// Messages are sharded by partition: cross-partition writes parallelize while +// ordering within a partition is preserved. +type Consumer struct { + reader MessageSource + sink Sink + logf func(format string, args ...interface{}) + workers int + shardBuf int + batchSize int + batchWait time.Duration + metrics *consumerMetrics +} + +const ( + sinkMaxAttempts = 5 + sinkBaseBackoff = 1 * time.Second + sinkMaxBackoff = 30 * time.Second + + defaultWorkers = 16 + defaultShardBuffer = 128 + defaultBatchSize = 16 + defaultBatchMaxWait = 10 * time.Millisecond +) + +// Backpressure: when the sink falls behind, ShardBufferSize fills, the fetcher +// blocks, and Kafka stops being polled. Zero values pick defaults. +type Options struct { + Logf func(format string, args ...interface{}) + Workers int + ShardBufferSize int + MaxBatchRecords int + BatchMaxWait time.Duration +} + +func New(reader MessageSource, sink Sink, opts Options) *Consumer { + logf := opts.Logf + if logf == nil { + logf = func(string, ...interface{}) {} + } + workers := opts.Workers + if workers <= 0 { + workers = defaultWorkers + } + shardBuf := opts.ShardBufferSize + if shardBuf <= 0 { + shardBuf = defaultShardBuffer + } + batchSize := opts.MaxBatchRecords + if batchSize <= 0 { + batchSize = defaultBatchSize + } + batchWait := opts.BatchMaxWait + if batchWait <= 0 { + batchWait = defaultBatchMaxWait + } + return &Consumer{ + reader: reader, + sink: sink, + logf: logf, + workers: workers, + shardBuf: shardBuf, + batchSize: batchSize, + batchWait: batchWait, + metrics: newConsumerMetrics(), + } +} + +// Run commits offsets only after the sink persists each message. +func (c *Consumer) Run(ctx context.Context) error { + g, gctx := errgroup.WithContext(ctx) + shards := make([]chan kafka.Message, c.workers) + for i := range shards { + shards[i] = make(chan kafka.Message, c.shardBuf) + ch := shards[i] + g.Go(func() error { return c.workerLoop(gctx, ch) }) + } + g.Go(func() error { + defer func() { + for _, ch := range shards { + close(ch) + } + }() + for { + msg, err := c.reader.FetchMessage(gctx) + if err != nil { + if isCancellation(err) { + return nil + } + return fmt.Errorf("fetch kafka message: %w", err) + } + shard := shardFor(msg.Partition, c.workers) + select { + case shards[shard] <- msg: + case <-gctx.Done(): + return nil + } + } + }) + if err := g.Wait(); err != nil && !isCancellation(err) { + return err + } + return nil +} + +func (c *Consumer) workerLoop(ctx context.Context, ch <-chan kafka.Message) error { + for { + select { + case <-ctx.Done(): + return nil + case msg, ok := <-ch: + if !ok { + return nil + } + msgs, ok := c.collectBatch(ctx, ch, msg) + if !ok { + return nil + } + if err := c.processBatch(ctx, msgs); err != nil { + if isCancellation(err) { + return nil + } + return err + } + } + } +} + +func (c *Consumer) collectBatch(ctx context.Context, ch <-chan kafka.Message, first kafka.Message) ([]kafka.Message, bool) { + msgs := make([]kafka.Message, 1, c.batchSize) + msgs[0] = first + if c.batchSize <= 1 { + return msgs, true + } + +drainBuffered: + for len(msgs) < c.batchSize { + select { + case <-ctx.Done(): + return nil, false + case msg, ok := <-ch: + if !ok { + return msgs, true + } + msgs = append(msgs, msg) + default: + break drainBuffered + } + } + if len(msgs) == c.batchSize { + return msgs, true + } + + timer := time.NewTimer(c.batchWait) + defer timer.Stop() + for len(msgs) < c.batchSize { + select { + case <-ctx.Done(): + return nil, false + case msg, ok := <-ch: + if !ok { + return msgs, true + } + msgs = append(msgs, msg) + case <-timer.C: + return msgs, true + } + } + return msgs, true +} + +func (c *Consumer) processBatch(ctx context.Context, msgs []kafka.Message) error { + records := make([]Record, 0, len(msgs)) + var firstVersion, lastVersion int64 + for i, msg := range msgs { + entry, err := DecodeEntry(msg.Value) + if err != nil { + return fmt.Errorf("decode message at offset %d: %w", msg.Offset, err) + } + if i == 0 { + firstVersion = entry.Version + } + lastVersion = entry.Version + records = append(records, Record{ + Topic: msg.Topic, + Partition: msg.Partition, + Offset: msg.Offset, + Entry: entry, + }) + } + start := time.Now() + if err := c.writeBatchWithRetry(ctx, records); err != nil { + return fmt.Errorf("sink write batch first_version=%d last_version=%d: %w", + firstVersion, lastVersion, err) + } + writeLatency := time.Since(start) + c.logf("wrote records=%d first_version=%d last_version=%d in %s", + len(records), firstVersion, lastVersion, writeLatency) + if err := c.reader.CommitMessages(ctx, msgs...); err != nil { + return fmt.Errorf("commit kafka offsets: %w", err) + } + c.metrics.recordBatch(ctx, int64(len(records)), batchLag(msgs), writeLatency) + return nil +} + +// batchLag reports how far behind the partition head the batch left the +// consumer: the max over messages of (high watermark - offset - 1). Returns -1 +// when no message carries a watermark so the gauge is left untouched. +func batchLag(msgs []kafka.Message) int64 { + lag := int64(-1) + for _, msg := range msgs { + if msg.HighWaterMark <= 0 { + continue + } + if l := msg.HighWaterMark - msg.Offset - 1; l > lag { + lag = l + } + } + return lag +} + +func (c *Consumer) writeBatchWithRetry(ctx context.Context, records []Record) error { + backoff := sinkBaseBackoff + var lastErr error + for attempt := 1; attempt <= sinkMaxAttempts; attempt++ { + err := c.sink.WriteBatch(ctx, records) + if err == nil { + return nil + } + lastErr = err + if isCancellation(err) { + return err + } + if attempt == sinkMaxAttempts { + break + } + c.logf("sink write attempt %d/%d failed: %v; retrying in %s", + attempt, sinkMaxAttempts, err, backoff) + if err := sleepWithContext(ctx, backoff); err != nil { + return err + } + backoff *= 2 + if backoff > sinkMaxBackoff { + backoff = sinkMaxBackoff + } + } + return fmt.Errorf("sink write failed after %d attempts: %w", sinkMaxAttempts, lastErr) +} + +func sleepWithContext(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func shardFor(partition, workers int) int { + if partition < 0 { + partition = -partition + } + return partition % workers +} + +func isCancellation(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) +} diff --git a/sei-db/state_db/ss/offload/consumer/kafka.go b/sei-db/state_db/ss/offload/consumer/kafka.go new file mode 100644 index 0000000000..894fbea0ce --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/kafka.go @@ -0,0 +1,126 @@ +package consumer + +import ( + "crypto/tls" + "fmt" + "strings" + "time" + + gogoproto "github.com/gogo/protobuf/proto" + "github.com/segmentio/kafka-go" + + dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload" +) + +// TLS/SASL must match the producer cluster. Commits are synchronous +// (kafka-go's zero CommitInterval) so offsets only advance after the sink +// persists each entry. +type KafkaReaderConfig struct { + Brokers []string + Topic string + GroupID string + ClientID string + Region string + StartOffset string // "first" or "last"; defaults to "first" + MinBytes int + MaxBytes int + MaxWait time.Duration + TLSEnabled bool + // SASLMechanism selects broker auth: "none", "plain" (username/password, + // e.g. Google Cloud Managed Kafka service-account credentials), or + // "aws-msk-iam". + SASLMechanism string + Username string + Password string +} + +func (c *KafkaReaderConfig) ApplyDefaults() { + if c.ClientID == "" { + c.ClientID = "sei-historical-offload-consumer" + } + if c.StartOffset == "" { + c.StartOffset = "first" + } + if c.MinBytes == 0 { + c.MinBytes = 1 + } + if c.MaxBytes == 0 { + c.MaxBytes = 10 << 20 + } + if c.MaxWait == 0 { + c.MaxWait = 500 * time.Millisecond + } +} + +func (c *KafkaReaderConfig) Validate() error { + if len(c.Brokers) == 0 { + return fmt.Errorf("kafka brokers are required") + } + if c.Topic == "" { + return fmt.Errorf("kafka topic is required") + } + if c.GroupID == "" { + return fmt.Errorf("kafka group id is required") + } + switch strings.ToLower(c.StartOffset) { + case "", "first", "last": + default: + return fmt.Errorf("unsupported kafka start offset %q", c.StartOffset) + } + return offload.ValidateSASL(c.saslConfig()) +} + +func (c KafkaReaderConfig) saslConfig() offload.KafkaConfig { + return offload.KafkaConfig{ + Region: c.Region, + TLSEnabled: c.TLSEnabled, + SASLMechanism: c.SASLMechanism, + Username: c.Username, + Password: c.Password, + } +} + +func NewKafkaReader(cfg KafkaReaderConfig) (*kafka.Reader, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + + dialer := &kafka.Dialer{ + ClientID: cfg.ClientID, + Timeout: 10 * time.Second, + } + if cfg.TLSEnabled { + dialer.TLS = &tls.Config{MinVersion: tls.VersionTLS12} + } + mech, err := offload.NewSASLMechanism(cfg.saslConfig()) + if err != nil { + return nil, err + } + dialer.SASLMechanism = mech + + start := kafka.FirstOffset + if strings.EqualFold(cfg.StartOffset, "last") { + start = kafka.LastOffset + } + + return kafka.NewReader(kafka.ReaderConfig{ + Brokers: cfg.Brokers, + Topic: cfg.Topic, + GroupID: cfg.GroupID, + Dialer: dialer, + MinBytes: cfg.MinBytes, + MaxBytes: cfg.MaxBytes, + MaxWait: cfg.MaxWait, + StartOffset: start, + }), nil +} + +func DecodeEntry(payload []byte) (*dbproto.ChangelogEntry, error) { + entry := &dbproto.ChangelogEntry{} + if err := gogoproto.Unmarshal(payload, entry); err != nil { + return nil, fmt.Errorf("decode changelog entry: %w", err) + } + return entry, nil +} diff --git a/sei-db/state_db/ss/offload/consumer/kafka_test.go b/sei-db/state_db/ss/offload/consumer/kafka_test.go new file mode 100644 index 0000000000..57056e9951 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/kafka_test.go @@ -0,0 +1,77 @@ +package consumer + +import ( + "testing" + + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/stretchr/testify/require" +) + +func TestKafkaReaderConfigApplyDefaults(t *testing.T) { + cfg := KafkaReaderConfig{ + Brokers: []string{"localhost:9092"}, + Topic: "historical-offload", + GroupID: "bigtable", + } + cfg.ApplyDefaults() + require.Equal(t, "sei-historical-offload-consumer", cfg.ClientID) + require.Equal(t, "first", cfg.StartOffset) + require.Equal(t, 1, cfg.MinBytes) + require.Equal(t, 10<<20, cfg.MaxBytes) +} + +func TestKafkaReaderConfigValidate(t *testing.T) { + cfg := KafkaReaderConfig{} + require.ErrorContains(t, cfg.Validate(), "brokers") + cfg = KafkaReaderConfig{Brokers: []string{"x"}} + require.ErrorContains(t, cfg.Validate(), "topic") + cfg = KafkaReaderConfig{Brokers: []string{"x"}, Topic: "t"} + require.ErrorContains(t, cfg.Validate(), "group id") + cfg = KafkaReaderConfig{ + Brokers: []string{"x"}, + Topic: "t", + GroupID: "g", + StartOffset: "middle", + } + require.ErrorContains(t, cfg.Validate(), "start offset") +} + +func TestKafkaReaderConfigValidateSASL(t *testing.T) { + base := KafkaReaderConfig{Brokers: []string{"x"}, Topic: "t", GroupID: "g"} + + cfg := base + cfg.SASLMechanism = "aws-msk-iam" + require.ErrorContains(t, cfg.Validate(), "tls") + + cfg.TLSEnabled = true + require.ErrorContains(t, cfg.Validate(), "region") + + cfg.Region = "us-east-1" + require.NoError(t, cfg.Validate()) + + // SASL/PLAIN (e.g. Google Cloud Managed Kafka) needs credentials. + cfg = base + cfg.SASLMechanism = "plain" + require.ErrorContains(t, cfg.Validate(), "username and password") + + cfg.Username = "svc@project.iam.gserviceaccount.com" + cfg.Password = "base64-encoded-key" + require.NoError(t, cfg.Validate()) + + cfg = base + cfg.SASLMechanism = "scram" + require.ErrorContains(t, cfg.Validate(), "sasl mechanism") + + cfg = base + cfg.SASLMechanism = "none" + require.NoError(t, cfg.Validate()) +} + +func TestDecodeEntry(t *testing.T) { + entry := &proto.ChangelogEntry{Version: 7} + payload, err := entry.Marshal() + require.NoError(t, err) + got, err := DecodeEntry(payload) + require.NoError(t, err) + require.Equal(t, int64(7), got.Version) +} diff --git a/sei-db/state_db/ss/offload/consumer/metrics.go b/sei-db/state_db/ss/offload/consumer/metrics.go new file mode 100644 index 0000000000..afcca00702 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/metrics.go @@ -0,0 +1,66 @@ +package consumer + +import ( + "context" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/metric" +) + +// consumerMeterName is the OpenTelemetry meter for the Kafka -> backend offload +// consumer. It exports through the same Prometheus pipeline as the sinks, so +// consumer throughput and lag sit alongside the backend cost counters. +const consumerMeterName = "seidb_offload_consumer" + +// consumerMetrics tracks Kafka drain throughput and lag. One consumer runs per +// process, so no attributes are needed — Prometheus job/instance labels +// separate backends. +type consumerMetrics struct { + recordsProcessed metric.Int64Counter + sinkWriteLatency metric.Float64Histogram + kafkaLag metric.Int64Gauge +} + +func newConsumerMetrics() *consumerMetrics { + meter := otel.Meter(consumerMeterName) + recordsProcessed, _ := meter.Int64Counter( + "consumer_records_processed_total", + metric.WithDescription("Changelog records written to the backend and committed to Kafka"), + metric.WithUnit("{record}"), + ) + sinkWriteLatency, _ := meter.Float64Histogram( + "consumer_sink_write_latency_seconds", + metric.WithDescription("Latency of a batch sink write including retries; the _count series is the batch count"), + metric.WithUnit("s"), + metric.WithExplicitBucketBoundaries( + 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, + ), + ) + kafkaLag, _ := meter.Int64Gauge( + "consumer_kafka_lag", + metric.WithDescription("Messages between the last processed offset and the partition high watermark (max across the batch)"), + metric.WithUnit("{message}"), + ) + return &consumerMetrics{ + recordsProcessed: recordsProcessed, + sinkWriteLatency: sinkWriteLatency, + kafkaLag: kafkaLag, + } +} + +// recordBatch reports one successfully written-and-committed batch: the number +// of records drained, how long the sink write took, and how far behind Kafka +// the batch left the consumer. +func (m *consumerMetrics) recordBatch(ctx context.Context, records int64, lag int64, writeLatency time.Duration) { + if m == nil { + return + } + if records > 0 { + m.recordsProcessed.Add(ctx, records) + } + m.sinkWriteLatency.Record(ctx, writeLatency.Seconds()) + if lag >= 0 { + m.kafkaLag.Record(ctx, lag) + } +} diff --git a/sei-db/state_db/ss/offload/consumer/metrics_test.go b/sei-db/state_db/ss/offload/consumer/metrics_test.go new file mode 100644 index 0000000000..7b77d3b4dc --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/metrics_test.go @@ -0,0 +1,36 @@ +package consumer + +import ( + "context" + "testing" + "time" + + "github.com/segmentio/kafka-go" + "github.com/stretchr/testify/require" +) + +func TestBatchLag(t *testing.T) { + // Max over the batch of (HighWaterMark - Offset - 1); the head message + // (offset == HWM-1) has lag 0. + msgs := []kafka.Message{ + {Offset: 10, HighWaterMark: 100}, // 89 + {Offset: 98, HighWaterMark: 100}, // 1 + {Offset: 99, HighWaterMark: 100}, // 0 (head) + } + require.Equal(t, int64(89), batchLag(msgs)) + + // No watermark information -> sentinel -1 so the gauge is left untouched. + require.Equal(t, int64(-1), batchLag([]kafka.Message{{Offset: 5}})) + require.Equal(t, int64(-1), batchLag(nil)) +} + +// recordBatch must be a safe no-op on a nil receiver and must not panic when +// the global MeterProvider is the default no-op provider. +func TestConsumerMetricsRecordNoPanic(t *testing.T) { + var nilM *consumerMetrics + nilM.recordBatch(context.Background(), 5, 3, time.Millisecond) + + m := newConsumerMetrics() + m.recordBatch(context.Background(), 128, 0, 12*time.Millisecond) + m.recordBatch(context.Background(), 0, -1, time.Millisecond) // zero records, unknown lag +} diff --git a/sei-db/state_db/ss/offload/consumer/sink.go b/sei-db/state_db/ss/offload/consumer/sink.go new file mode 100644 index 0000000000..deef46c336 --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/sink.go @@ -0,0 +1,21 @@ +package consumer + +import ( + "context" + + dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" +) + +// Topic/Partition/Offset are kept alongside Entry so sinks can be idempotent +// across replayed Kafka messages. +type Record struct { + Topic string + Partition int + Offset int64 + Entry *dbproto.ChangelogEntry +} + +type Sink interface { + WriteBatch(ctx context.Context, records []Record) error + Close() error +} diff --git a/sei-db/state_db/ss/offload/historical/bigtable.go b/sei-db/state_db/ss/offload/historical/bigtable.go new file mode 100644 index 0000000000..edb85510fa --- /dev/null +++ b/sei-db/state_db/ss/offload/historical/bigtable.go @@ -0,0 +1,749 @@ +package historical + +import ( + "context" + "encoding/binary" + "fmt" + "hash/fnv" + "io" + "regexp" + "strings" + "time" + + "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "golang.org/x/sync/errgroup" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/credentials/oauth" + "google.golang.org/grpc/keepalive" + "google.golang.org/grpc/status" +) + +const ( + DefaultBigtableFamily = "state" + DefaultBigtableShards = 256 + + bigtableMutationPrefix = byte('m') + bigtableVersionPrefix = byte('v') + bigtableUpgradePrefix = byte('u') + + BigtableValueColumn = "value" + BigtableDeletedColumn = "deleted" + + // VersionBucketCount spreads monotonically increasing block-version markers + // across a bounded set of row prefixes while keeping LastVersion cheap. + VersionBucketCount = 64 +) + +// VersionBucket maps a version to its marker bucket. +func VersionBucket(version int64) int { + if version < 0 { + version = -version + } + return int(version % VersionBucketCount) +} + +const ( + bigtableEndpoint = "bigtable.googleapis.com:443" + + defaultBigtableReadWorkers = 16 + + // Transient stream failures (tablet moves, GFE restarts) surface as + // UNAVAILABLE; retry reads a bounded number of times with backoff. + bigtableReadAttempts = 3 + bigtableReadRetryBackoff = 100 * time.Millisecond + + maxUint16Int = 1<<16 - 1 + maxUint32Int = 1<<32 - 1 + maxInt64Uint64 = 1<<63 - 1 +) + +// BigtableClient speaks the Bigtable data protocol directly over gRPC. The +// official cloud.google.com/go/bigtable client requires a newer +// google.golang.org/grpc than the repo-wide v1.57 replace pin allows, so this +// package carries its own thin transport: ReadRows chunk assembly, MutateRows +// result accounting, bounded read retries, and per-RPC cost metrics. +type BigtableClient struct { + conn *grpc.ClientConn + data bigtablepb.BigtableClient + tableName string + appProfile string + // table is the short table name used as the metric attribute. + table string + metrics *bigtableMetrics +} + +type BigtableCell struct { + Family string + Qualifier string + Value []byte +} + +type BigtableRow struct { + Key string + Cells []BigtableCell +} + +type BigtableSetCell struct { + Family string + Qualifier string + TimestampMicros int64 + Value []byte +} + +type BigtableRowMutation struct { + RowKey string + SetCells []BigtableSetCell +} + +type BigtableConfig struct { + ProjectID string + InstanceID string + Table string + Family string + AppProfile string + Shards int +} + +func (c *BigtableConfig) ApplyDefaults() { + if c.Family == "" { + c.Family = DefaultBigtableFamily + } + if c.Shards == 0 { + c.Shards = DefaultBigtableShards + } +} + +func (c BigtableConfig) Configured() bool { + return strings.TrimSpace(c.ProjectID) != "" || + strings.TrimSpace(c.InstanceID) != "" || + strings.TrimSpace(c.Table) != "" +} + +func (c *BigtableConfig) Validate() error { + if strings.TrimSpace(c.ProjectID) == "" { + return fmt.Errorf("bigtable project id is required") + } + if strings.TrimSpace(c.InstanceID) == "" { + return fmt.Errorf("bigtable instance id is required") + } + if strings.TrimSpace(c.Table) == "" { + return fmt.Errorf("bigtable table is required") + } + if strings.TrimSpace(c.Family) == "" { + return fmt.Errorf("bigtable family is required") + } + if c.Shards <= 0 || c.Shards > maxUint16Int { + return fmt.Errorf("bigtable shards must be between 1 and 65535") + } + return nil +} + +func NewBigtableReader(cfg BigtableConfig) (Reader, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + client, err := OpenBigtableClient(context.Background(), cfg) + if err != nil { + return nil, err + } + return &bigtableReader{ + client: client, + readRows: client.ReadRows, + family: cfg.Family, + shards: cfg.Shards, + }, nil +} + +func OpenBigtableClient(ctx context.Context, cfg BigtableConfig) (*BigtableClient, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + perRPC, err := oauth.NewApplicationDefault(ctx, "https://www.googleapis.com/auth/bigtable.data", "https://www.googleapis.com/auth/cloud-platform") + if err != nil { + return nil, fmt.Errorf("bigtable auth: %w", err) + } + conn, err := grpc.DialContext(ctx, bigtableEndpoint, + grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), + grpc.WithPerRPCCredentials(perRPC), + // Detect silently dropped connections instead of hanging reads on the + // OS TCP timeout. + grpc.WithKeepaliveParams(keepalive.ClientParameters{ + Time: 30 * time.Second, + Timeout: 10 * time.Second, + }), + ) + if err != nil { + return nil, fmt.Errorf("open bigtable connection: %w", err) + } + return &BigtableClient{ + conn: conn, + data: bigtablepb.NewBigtableClient(conn), + tableName: bigtableTableName(cfg.ProjectID, cfg.InstanceID, cfg.Table), + appProfile: cfg.AppProfile, + table: cfg.Table, + metrics: newBigtableMetrics(), + }, nil +} + +type BigtableReadRowsFunc func(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error + +type BigtableApplyBulkFunc func(ctx context.Context, rows []BigtableRowMutation) ([]error, error) + +func (c *BigtableClient) Close() error { + if c.conn == nil { + return nil + } + return c.conn.Close() +} + +// ReadRows retries transient stream failures as long as no row has been +// delivered to the callback yet, which covers every limit-1 point read and +// version-bucket scan in this package. +func (c *BigtableClient) ReadRows(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + return readRowsWithRetry(ctx, c.readRowsOnce, startKey, endKey, limit, family, f, qualifiers...) +} + +func readRowsWithRetry(ctx context.Context, once BigtableReadRowsFunc, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + backoff := bigtableReadRetryBackoff + var delivered bool + for attempt := 1; ; attempt++ { + err := once(ctx, startKey, endKey, limit, family, func(row BigtableRow) bool { + delivered = true + return f(row) + }, qualifiers...) + if err == nil || delivered || attempt >= bigtableReadAttempts || status.Code(err) != codes.Unavailable { + return err + } + if sleepErr := sleepWithContext(ctx, backoff); sleepErr != nil { + return err + } + backoff *= 2 + } +} + +func sleepWithContext(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *BigtableClient) readRowsOnce(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + start := time.Now() + var rowsRead, bytesRead int64 + defer func() { + c.metrics.recordRead(ctx, c.table, time.Since(start), rowsRead, bytesRead) + }() + req := &bigtablepb.ReadRowsRequest{ + TableName: c.tableName, + AppProfileId: c.appProfile, + Rows: &bigtablepb.RowSet{RowRanges: []*bigtablepb.RowRange{{ + StartKey: &bigtablepb.RowRange_StartKeyClosed{StartKeyClosed: startKey}, + EndKey: &bigtablepb.RowRange_EndKeyOpen{EndKeyOpen: endKey}, + }}}, + RowsLimit: limit, + } + if len(endKey) == 0 { + req.Rows.RowRanges[0].EndKey = nil + } + req.Filter = bigtableReadFilter(family, qualifiers...) + // Cancel on return so an early callback exit tears down the server stream + // instead of leaving it to the connection finalizer. + streamCtx, cancel := context.WithCancel(ctx) + defer cancel() + stream, err := c.data.ReadRows(streamCtx, req) + if err != nil { + return err + } + var builder bigtableRowBuilder + for { + resp, err := stream.Recv() + if err == io.EOF { + return nil + } + if err != nil { + return err + } + for _, chunk := range resp.Chunks { + row, committed, err := builder.add(chunk) + if err != nil { + return err + } + if committed { + rowsRead++ + bytesRead += bigtableRowSize(row) + if !f(row) { + return nil + } + } + } + } +} + +func (c *BigtableClient) ApplyBulk(ctx context.Context, rows []BigtableRowMutation) ([]error, error) { + if len(rows) == 0 { + return nil, nil + } + start := time.Now() + var bytesWritten int64 + defer func() { + c.metrics.recordWrite(ctx, c.table, time.Since(start), int64(len(rows)), bytesWritten) + }() + entries := make([]*bigtablepb.MutateRowsRequest_Entry, 0, len(rows)) + for _, row := range rows { + bytesWritten += bigtableMutationSize(row) + entry := &bigtablepb.MutateRowsRequest_Entry{ + RowKey: []byte(row.RowKey), + Mutations: make([]*bigtablepb.Mutation, 0, len(row.SetCells)), + } + for _, cell := range row.SetCells { + entry.Mutations = append(entry.Mutations, &bigtablepb.Mutation{ + Mutation: &bigtablepb.Mutation_SetCell_{SetCell: &bigtablepb.Mutation_SetCell{ + FamilyName: cell.Family, + ColumnQualifier: []byte(cell.Qualifier), + TimestampMicros: cell.TimestampMicros, + Value: cell.Value, + }}, + }) + } + entries = append(entries, entry) + } + stream, err := c.data.MutateRows(ctx, &bigtablepb.MutateRowsRequest{ + TableName: c.tableName, + AppProfileId: c.appProfile, + Entries: entries, + }) + if err != nil { + return nil, err + } + rowErrs := make([]error, len(rows)) + seen := make([]bool, len(rows)) + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return rowErrs, err + } + for _, entry := range resp.Entries { + if entry.Index < 0 || int(entry.Index) >= len(rowErrs) { + return rowErrs, fmt.Errorf("bigtable returned invalid mutation index %d", entry.Index) + } + idx := int(entry.Index) + seen[idx] = true + if st := entry.Status; st != nil && st.Code != 0 { + rowErrs[idx] = fmt.Errorf("bigtable status %d: %s", st.Code, st.Message) + } + } + } + for i := range seen { + if !seen[i] { + rowErrs[i] = fmt.Errorf("bigtable missing mutation result") + } + } + return rowErrs, nil +} + +type bigtableReader struct { + client *BigtableClient + readRows BigtableReadRowsFunc + family string + shards int +} + +var _ Reader = (*bigtableReader)(nil) + +func (r *bigtableReader) Close() error { + if r.client != nil { + return r.client.Close() + } + return nil +} + +func (r *bigtableReader) LastVersion(ctx context.Context) (int64, error) { + return BigtableLastVersion(ctx, r.readRows) +} + +func (r *bigtableReader) Has(ctx context.Context, storeName string, key []byte, targetVersion int64) (bool, error) { + start, end := bigtableMutationRowRange(storeName, key, targetVersion, r.shards) + var row BigtableRow + err := r.readRows(ctx, start, end, 1, r.family, func(r BigtableRow) bool { + row = r + return false + }, BigtableDeletedColumn) + if err != nil { + return false, fmt.Errorf("bigtable has lookup: %w", err) + } + if row.Key == "" { + return false, nil + } + deleted, err := bigtableDeletedFromRow(row, r.family) + if err != nil { + return false, err + } + return !deleted, nil +} + +func (r *bigtableReader) Get(ctx context.Context, storeName string, key []byte, targetVersion int64) (Value, error) { + start, end := bigtableMutationRowRange(storeName, key, targetVersion, r.shards) + var row BigtableRow + err := r.readRows(ctx, start, end, 1, r.family, func(r BigtableRow) bool { + row = r + return false + }, BigtableValueColumn, BigtableDeletedColumn) + if err != nil { + return Value{}, fmt.Errorf("bigtable get lookup: %w", err) + } + if row.Key == "" { + return Value{}, ErrNotFound + } + return bigtableValueFromRow(row, r.family) +} + +func BigtableLastVersion(ctx context.Context, readRows BigtableReadRowsFunc) (int64, error) { + versions := make([]int64, VersionBucketCount) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(defaultBigtableReadWorkers) + for bucket := 0; bucket < VersionBucketCount; bucket++ { + bucket := bucket + g.Go(func() error { + prefix := bigtableVersionRowPrefix(bucket) + var bucketVersion int64 + err := readRows(gctx, prefix, bigtablePrefixEnd(prefix), 1, "", func(row BigtableRow) bool { + if version, ok := bigtableVersionFromRowKey(row.Key); ok { + bucketVersion = version + } + return false + }) + if err != nil { + return fmt.Errorf("read latest bigtable version bucket %d: %w", bucket, err) + } + versions[bucket] = bucketVersion + return nil + }) + } + if err := g.Wait(); err != nil { + return 0, err + } + var maxVersion int64 + for _, version := range versions { + if version > maxVersion { + maxVersion = version + } + } + return maxVersion, nil +} + +// bigtableValueFromRow interprets a mutation row. The returned Value aliases +// the row's cell buffer, which the row builder allocates per cell. +func bigtableValueFromRow(row BigtableRow, family string) (Value, error) { + version, ok := bigtableVersionFromRowKey(row.Key) + if !ok { + return Value{}, fmt.Errorf("invalid bigtable mutation row key") + } + var value []byte + deleted := false + for _, cell := range row.Cells { + if cell.Family != family { + continue + } + switch cell.Qualifier { + case BigtableValueColumn: + value = cell.Value + case BigtableDeletedColumn: + deleted = len(cell.Value) > 0 && cell.Value[0] == 1 + } + } + if deleted || value == nil { + return Value{}, ErrNotFound + } + return Value{Bytes: value, Version: version}, nil +} + +func bigtableDeletedFromRow(row BigtableRow, family string) (bool, error) { + if _, ok := bigtableVersionFromRowKey(row.Key); !ok { + return false, fmt.Errorf("invalid bigtable mutation row key") + } + for _, cell := range row.Cells { + if cell.Family == family && cell.Qualifier == BigtableDeletedColumn { + return len(cell.Value) > 0 && cell.Value[0] == 1, nil + } + } + return false, nil +} + +func bigtableMutationRowPrefixBytes(storeName string, key []byte, shards int) []byte { + shards = normalizeBigtableShards(shards) + shard := bigtableShard(storeName, key, shards) + prefixLen := 1 + 2 + 2 + len(storeName) + 4 + len(key) + prefix := make([]byte, prefixLen, prefixLen+8) + prefix[0] = bigtableMutationPrefix + binary.BigEndian.PutUint16(prefix[1:], shard) + binary.BigEndian.PutUint16(prefix[3:], uint16FromBoundedInt(len(storeName))) + copy(prefix[5:], storeName) + keyOffset := 5 + len(storeName) + binary.BigEndian.PutUint32(prefix[keyOffset:], uint32FromBoundedInt(len(key))) + copy(prefix[keyOffset+4:], key) + return prefix +} + +func BigtableMutationRowKey(storeName string, key []byte, version int64, shards int) string { + return string(bigtableMutationRowKeyBytes(storeName, key, version, shards)) +} + +func bigtableMutationRowKeyBytes(storeName string, key []byte, version int64, shards int) []byte { + prefix := bigtableMutationRowPrefixBytes(storeName, key, shards) + return append(prefix, bigtableInvertedVersion(version)...) +} + +// bigtableMutationRowRange returns the start key of a point lookup at +// targetVersion and the exclusive end of the key's version range, hashing and +// encoding the shared row prefix only once. +func bigtableMutationRowRange(storeName string, key []byte, targetVersion int64, shards int) (start, end []byte) { + prefix := bigtableMutationRowPrefixBytes(storeName, key, shards) + end = bigtablePrefixEnd(prefix) + start = append(prefix, bigtableInvertedVersion(targetVersion)...) + return start, end +} + +func BigtableVersionRowKey(version int64) string { + prefix := bigtableVersionRowPrefix(VersionBucket(version)) + return string(append(prefix, bigtableInvertedVersion(version)...)) +} + +func BigtableUpgradeRowKey(version int64, name string) string { + key := make([]byte, 1+8+2+len(name)) + key[0] = bigtableUpgradePrefix + copy(key[1:], bigtableInvertedVersion(version)) + binary.BigEndian.PutUint16(key[9:], uint16FromBoundedInt(len(name))) + copy(key[11:], name) + return string(key) +} + +func bigtableVersionFromRowKey(rowKey string) (int64, bool) { + key := []byte(rowKey) + switch { + case len(key) >= 1+2+8 && key[0] == bigtableVersionPrefix: + return bigtableDecodeInvertedVersion(key[3:11]) + case len(key) >= 8 && key[0] == bigtableMutationPrefix: + return bigtableDecodeInvertedVersion(key[len(key)-8:]) + default: + return 0, false + } +} + +func BigtableTimestamp(version int64) int64 { + return version * int64(time.Millisecond/time.Microsecond) +} + +type bigtableRowBuilder struct { + key []byte + family string + qualifier string + value []byte + inCell bool + cells []BigtableCell +} + +func (b *bigtableRowBuilder) add(chunk *bigtablepb.ReadRowsResponse_CellChunk) (BigtableRow, bool, error) { + if chunk.GetResetRow() { + b.reset() + return BigtableRow{}, false, nil + } + if len(chunk.RowKey) != 0 { + b.key = append(b.key[:0], chunk.RowKey...) + b.family = "" + b.qualifier = "" + b.value = nil + b.inCell = false + b.cells = b.cells[:0] + } + if chunk.FamilyName != nil { + b.family = chunk.FamilyName.Value + } + if chunk.Qualifier != nil { + b.qualifier = string(chunk.Qualifier.Value) + b.value = b.value[:0] + b.inCell = true + } + // A non-final chunk of a split cell advertises the total value size; grow + // the buffer once instead of re-growing on every continuation chunk. + if size := int(chunk.ValueSize); size > cap(b.value) { + grown := make([]byte, len(b.value), size) + copy(grown, b.value) + b.value = grown + } + b.value = append(b.value, chunk.Value...) + if b.inCell && chunk.ValueSize == 0 { + b.cells = append(b.cells, BigtableCell{ + Family: b.family, + Qualifier: b.qualifier, + Value: append([]byte(nil), b.value...), + }) + b.value = nil + b.inCell = false + } + if !chunk.GetCommitRow() { + return BigtableRow{}, false, nil + } + if len(b.key) == 0 { + return BigtableRow{}, false, fmt.Errorf("bigtable committed row without key") + } + row := BigtableRow{ + Key: string(append([]byte(nil), b.key...)), + Cells: append([]BigtableCell(nil), b.cells...), + } + b.reset() + return row, true, nil +} + +func (b *bigtableRowBuilder) reset() { + b.key = nil + b.family = "" + b.qualifier = "" + b.value = nil + b.inCell = false + b.cells = nil +} + +func bigtableTableName(projectID, instanceID, table string) string { + return fmt.Sprintf("projects/%s/instances/%s/tables/%s", projectID, instanceID, table) +} + +func bigtableReadFilter(family string, qualifiers ...string) *bigtablepb.RowFilter { + filters := make([]*bigtablepb.RowFilter, 0, 2) + if family != "" { + filters = append(filters, &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_FamilyNameRegexFilter{FamilyNameRegexFilter: regexp.QuoteMeta(family)}, + }) + } + if qualifierFilter := bigtableQualifierFilter(qualifiers...); qualifierFilter != nil { + filters = append(filters, qualifierFilter) + } + switch len(filters) { + case 0: + return nil + case 1: + return filters[0] + default: + return &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_Chain_{Chain: &bigtablepb.RowFilter_Chain{Filters: filters}}, + } + } +} + +func bigtableQualifierFilter(qualifiers ...string) *bigtablepb.RowFilter { + filters := make([]*bigtablepb.RowFilter, 0, len(qualifiers)) + for _, qualifier := range qualifiers { + if qualifier == "" { + continue + } + filters = append(filters, &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_ColumnQualifierRegexFilter{ + ColumnQualifierRegexFilter: []byte(regexp.QuoteMeta(qualifier)), + }, + }) + } + switch len(filters) { + case 0: + return nil + case 1: + return filters[0] + default: + return &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_Interleave_{Interleave: &bigtablepb.RowFilter_Interleave{Filters: filters}}, + } + } +} + +func bigtableVersionRowPrefix(bucket int) []byte { + prefix := make([]byte, 1+2) + prefix[0] = bigtableVersionPrefix + binary.BigEndian.PutUint16(prefix[1:], uint16FromBoundedInt(bucket)) + return prefix +} + +func bigtablePrefixEnd(prefix []byte) []byte { + end := append([]byte(nil), prefix...) + for i := len(end) - 1; i >= 0; i-- { + if end[i] != 0xff { + end[i]++ + return end[:i+1] + } + } + return nil +} + +func bigtableInvertedVersion(version int64) []byte { + out := make([]byte, 8) + binary.BigEndian.PutUint64(out, ^uint64FromNonNegativeInt64(version)) + return out +} + +func bigtableDecodeInvertedVersion(encoded []byte) (int64, bool) { + version := ^binary.BigEndian.Uint64(encoded) + if version > maxInt64Uint64 { + return 0, false + } + // #nosec G115 -- version is checked above to fit in int64. + return int64(version), true +} + +func bigtableShard(storeName string, key []byte, shards int) uint16 { + h := fnv.New32a() + _, _ = h.Write([]byte(storeName)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write(key) + return uint16FromBoundedUint32(h.Sum32() % uint32FromBoundedInt(shards)) +} + +func normalizeBigtableShards(shards int) int { + if shards <= 0 { + return DefaultBigtableShards + } + if shards > maxUint16Int { + return maxUint16Int + } + return shards +} + +func uint16FromBoundedInt(value int) uint16 { + if value < 0 || value > maxUint16Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) + } + // #nosec G115 -- value is checked above to fit in uint16. + return uint16(value) +} + +func uint32FromBoundedInt(value int) uint32 { + if value < 0 || value > maxUint32Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint32", value)) + } + // #nosec G115 -- value is checked above to fit in uint32. + return uint32(value) +} + +func uint16FromBoundedUint32(value uint32) uint16 { + if value > maxUint16Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) + } + // #nosec G115 -- value is checked above to fit in uint16. + return uint16(value) +} + +func uint64FromNonNegativeInt64(value int64) uint64 { + if value < 0 { + panic(fmt.Sprintf("bigtable version %d is negative", value)) + } + // #nosec G115 -- value is checked above to be non-negative. + return uint64(value) +} diff --git a/sei-db/state_db/ss/offload/historical/bigtable_test.go b/sei-db/state_db/ss/offload/historical/bigtable_test.go new file mode 100644 index 0000000000..e5a93e9291 --- /dev/null +++ b/sei-db/state_db/ss/offload/historical/bigtable_test.go @@ -0,0 +1,242 @@ +package historical + +import ( + "context" + "fmt" + "sort" + "sync" + "testing" + + "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/wrapperspb" +) + +func TestBigtableConfigDefaultsAndValidate(t *testing.T) { + cfg := BigtableConfig{ + ProjectID: "project", + InstanceID: "instance", + Table: "state", + } + cfg.ApplyDefaults() + require.Equal(t, DefaultBigtableFamily, cfg.Family) + require.Equal(t, DefaultBigtableShards, cfg.Shards) + require.NoError(t, cfg.Validate()) + + missingProject := BigtableConfig{InstanceID: "i", Table: "t", Family: "f"} + missingInstance := BigtableConfig{ProjectID: "p", Table: "t", Family: "f"} + missingTable := BigtableConfig{ProjectID: "p", InstanceID: "i", Family: "f"} + missingShards := BigtableConfig{ProjectID: "p", InstanceID: "i", Table: "t", Family: "f"} + require.ErrorContains(t, missingProject.Validate(), "project") + require.ErrorContains(t, missingInstance.Validate(), "instance") + require.ErrorContains(t, missingTable.Validate(), "table") + require.ErrorContains(t, missingShards.Validate(), "shards") +} + +func TestBigtableMutationRowKeyOrdersLatestVersionFirst(t *testing.T) { + key40 := BigtableMutationRowKey("bank", []byte("k1"), 40, 256) + key60 := BigtableMutationRowKey("bank", []byte("k1"), 60, 256) + key80 := BigtableMutationRowKey("bank", []byte("k1"), 80, 256) + keys := []string{key40, key80, key60} + sort.Strings(keys) + require.Equal(t, []string{key80, key60, key40}, keys) + + version, ok := bigtableVersionFromRowKey(key60) + require.True(t, ok) + require.Equal(t, int64(60), version) + require.NotEqual(t, + bigtableMutationRowPrefixBytes("bank", []byte("k"), 256), + bigtableMutationRowPrefixBytes("bank", []byte("k1"), 256), + ) +} + +func TestBigtableValueFromRow(t *testing.T) { + rowKey := BigtableMutationRowKey("bank", []byte("k"), 7, 256) + row := BigtableRow{ + Key: rowKey, + Cells: []BigtableCell{ + {Family: DefaultBigtableFamily, Qualifier: BigtableValueColumn, Value: []byte("value")}, + {Family: DefaultBigtableFamily, Qualifier: BigtableDeletedColumn, Value: []byte{0}}, + }, + } + value, err := bigtableValueFromRow(row, DefaultBigtableFamily) + require.NoError(t, err) + require.Equal(t, []byte("value"), value.Bytes) + require.Equal(t, int64(7), value.Version) + + row.Cells[1].Value = []byte{1} + _, err = bigtableValueFromRow(row, DefaultBigtableFamily) + require.ErrorIs(t, err, ErrNotFound) +} + +func TestBigtableReaderGetUsesMVCCRange(t *testing.T) { + wantRow := BigtableMutationRowKey("bank", []byte("k"), 40, 256) + reader := &bigtableReader{ + family: DefaultBigtableFamily, + shards: 256, + readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + require.Equal(t, []byte(BigtableMutationRowKey("bank", []byte("k"), 60, 256)), startKey) + require.NotEmpty(t, endKey) + require.Equal(t, int64(1), limit) + require.Equal(t, DefaultBigtableFamily, family) + require.Equal(t, []string{BigtableValueColumn, BigtableDeletedColumn}, qualifiers) + f(BigtableRow{ + Key: wantRow, + Cells: []BigtableCell{ + {Family: DefaultBigtableFamily, Qualifier: BigtableValueColumn, Value: []byte("v40")}, + {Family: DefaultBigtableFamily, Qualifier: BigtableDeletedColumn, Value: []byte{0}}, + }, + }) + return nil + }, + } + + value, err := reader.Get(context.Background(), "bank", []byte("k"), 60) + require.NoError(t, err) + require.Equal(t, []byte("v40"), value.Bytes) + require.Equal(t, int64(40), value.Version) +} + +func TestBigtableReaderHasOnlyReadsDeletedColumn(t *testing.T) { + wantRow := BigtableMutationRowKey("bank", []byte("k"), 40, 256) + reader := &bigtableReader{ + family: DefaultBigtableFamily, + shards: 256, + readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + require.Equal(t, []byte(BigtableMutationRowKey("bank", []byte("k"), 60, 256)), startKey) + require.NotEmpty(t, endKey) + require.Equal(t, int64(1), limit) + require.Equal(t, DefaultBigtableFamily, family) + require.Equal(t, []string{BigtableDeletedColumn}, qualifiers) + f(BigtableRow{ + Key: wantRow, + Cells: []BigtableCell{{ + Family: DefaultBigtableFamily, + Qualifier: BigtableDeletedColumn, + Value: []byte{0}, + }}, + }) + return nil + }, + } + + ok, err := reader.Has(context.Background(), "bank", []byte("k"), 60) + require.NoError(t, err) + require.True(t, ok) +} + +// A cell split across chunks carries the total size in every chunk except the +// last, whose ValueSize is zero; only then may the cell be committed. +func TestBigtableRowBuilderAssemblesSplitCells(t *testing.T) { + var b bigtableRowBuilder + + row, committed, err := b.add(&bigtablepb.ReadRowsResponse_CellChunk{ + RowKey: []byte("rk"), + FamilyName: wrapperspb.String(DefaultBigtableFamily), + Qualifier: wrapperspb.Bytes([]byte(BigtableValueColumn)), + Value: []byte("hel"), + ValueSize: 8, + }) + require.NoError(t, err) + require.False(t, committed) + require.Empty(t, row.Key) + + _, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{Value: []byte("lo "), ValueSize: 8}) + require.NoError(t, err) + require.False(t, committed) + + row, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{ + Value: []byte("bt"), + RowStatus: &bigtablepb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true}, + }) + require.NoError(t, err) + require.True(t, committed) + require.Equal(t, "rk", row.Key) + require.Len(t, row.Cells, 1) + require.Equal(t, BigtableValueColumn, row.Cells[0].Qualifier) + require.Equal(t, []byte("hello bt"), row.Cells[0].Value) +} + +func TestReadRowsWithRetryRetriesUnavailable(t *testing.T) { + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(BigtableRow) bool, _ ...string) error { + attempts++ + if attempts < 3 { + return status.Error(codes.Unavailable, "tablet moved") + } + f(BigtableRow{Key: "rk"}) + return nil + } + var got []string + err := readRowsWithRetry(context.Background(), once, []byte("a"), []byte("b"), 1, "state", func(row BigtableRow) bool { + got = append(got, row.Key) + return false + }) + require.NoError(t, err) + require.Equal(t, 3, attempts) + require.Equal(t, []string{"rk"}, got) +} + +func TestReadRowsWithRetryStopsAfterDelivery(t *testing.T) { + // Once a row reached the callback, a retry would re-deliver it; the error + // must surface instead. + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(BigtableRow) bool, _ ...string) error { + attempts++ + f(BigtableRow{Key: "rk"}) + return status.Error(codes.Unavailable, "mid-stream drop") + } + err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(BigtableRow) bool { return true }) + require.Error(t, err) + require.Equal(t, 1, attempts) +} + +func TestReadRowsWithRetryDoesNotRetryNonRetryable(t *testing.T) { + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, _ func(BigtableRow) bool, _ ...string) error { + attempts++ + return status.Error(codes.NotFound, "table missing") + } + err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(BigtableRow) bool { return true }) + require.Error(t, err) + require.Equal(t, 1, attempts) +} + +func TestVersionBucket(t *testing.T) { + require.Equal(t, 0, VersionBucket(0)) + require.Equal(t, 1, VersionBucket(1)) + require.Equal(t, 0, VersionBucket(VersionBucketCount)) + require.Equal(t, 1, VersionBucket(-1)) +} + +func TestBigtableLastVersionScansBuckets(t *testing.T) { + versions := map[int]int64{ + 3: 42, + 9: 70, + } + seen := make(map[int]struct{}, VersionBucketCount) + var mu sync.Mutex + + got, err := BigtableLastVersion(context.Background(), func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { + if len(startKey) != 3 || startKey[0] != bigtableVersionPrefix { + return fmt.Errorf("unexpected start key %q", startKey) + } + if len(endKey) == 0 || limit != 1 || family != "" || len(qualifiers) != 0 { + return fmt.Errorf("unexpected scan params") + } + bucket := int(startKey[1])<<8 | int(startKey[2]) + mu.Lock() + seen[bucket] = struct{}{} + version := versions[bucket] + mu.Unlock() + if version > 0 { + f(BigtableRow{Key: BigtableVersionRowKey(version)}) + } + return nil + }) + require.NoError(t, err) + require.Equal(t, int64(70), got) + require.Len(t, seen, VersionBucketCount) +} diff --git a/sei-db/state_db/ss/offload/historical/metrics.go b/sei-db/state_db/ss/offload/historical/metrics.go new file mode 100644 index 0000000000..a542121fe1 --- /dev/null +++ b/sei-db/state_db/ss/offload/historical/metrics.go @@ -0,0 +1,135 @@ +package historical + +import ( + "context" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// bigtableMeterName is the OpenTelemetry meter used for Bigtable-backed +// historical-state access. It exports through the same process-wide +// MeterProvider (Prometheus exporter) as the rest of SeiDB. +const bigtableMeterName = "seidb_bigtable" + +// bigtableMetrics measures the cost of Bigtable-backed historical state so it +// is easy to see how read and write volume scales with historical query and +// ingestion load. The three cost drivers Bigtable bills on — request count, +// rows touched, and bytes transferred — are each captured directly: +// +// - request count: the *_latency_seconds histogram's _count series (free); +// - rows touched: bigtable_rows_read_total / bigtable_rows_mutated_total; +// - bytes moved: bigtable_bytes_read_total / bigtable_bytes_written_total. +// +// All series carry only a single low-cardinality "table" attribute, so the +// storage footprint stays flat as request volume grows. +type bigtableMetrics struct { + readLatency metric.Float64Histogram + rowsRead metric.Int64Counter + bytesRead metric.Int64Counter + writeLatency metric.Float64Histogram + rowsMutated metric.Int64Counter + bytesWritten metric.Int64Counter +} + +func newBigtableMetrics() *bigtableMetrics { + meter := otel.Meter(bigtableMeterName) + // Buckets span single-digit-ms cache-warm reads to multi-second slow RPCs. + latencyBuckets := metric.WithExplicitBucketBoundaries( + 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, + ) + readLatency, _ := meter.Float64Histogram( + "bigtable_read_latency_seconds", + metric.WithDescription("Latency of Bigtable ReadRows RPCs for historical state; the _count series doubles as the read request count"), + metric.WithUnit("s"), + latencyBuckets, + ) + rowsRead, _ := meter.Int64Counter( + "bigtable_rows_read_total", + metric.WithDescription("Rows returned by Bigtable ReadRows RPCs for historical state"), + metric.WithUnit("{row}"), + ) + bytesRead, _ := meter.Int64Counter( + "bigtable_bytes_read_total", + metric.WithDescription("Bytes (row keys + cell values) returned by Bigtable reads for historical state"), + metric.WithUnit("By"), + ) + writeLatency, _ := meter.Float64Histogram( + "bigtable_mutate_latency_seconds", + metric.WithDescription("Latency of Bigtable MutateRows RPCs for historical-state ingestion; the _count series doubles as the mutate request count"), + metric.WithUnit("s"), + latencyBuckets, + ) + rowsMutated, _ := meter.Int64Counter( + "bigtable_rows_mutated_total", + metric.WithDescription("Rows written by Bigtable MutateRows RPCs for historical-state ingestion"), + metric.WithUnit("{row}"), + ) + bytesWritten, _ := meter.Int64Counter( + "bigtable_bytes_written_total", + metric.WithDescription("Bytes (row keys + cell values) written by Bigtable mutations for historical state"), + metric.WithUnit("By"), + ) + return &bigtableMetrics{ + readLatency: readLatency, + rowsRead: rowsRead, + bytesRead: bytesRead, + writeLatency: writeLatency, + rowsMutated: rowsMutated, + bytesWritten: bytesWritten, + } +} + +// recordRead reports one ReadRows RPC: its latency plus the rows and bytes it +// returned. It records even when the RPC fails, since a failed read still costs. +func (m *bigtableMetrics) recordRead(ctx context.Context, table string, latency time.Duration, rows, bytes int64) { + if m == nil { + return + } + attrs := metric.WithAttributes(attribute.String("table", table)) + m.readLatency.Record(ctx, latency.Seconds(), attrs) + if rows > 0 { + m.rowsRead.Add(ctx, rows, attrs) + } + if bytes > 0 { + m.bytesRead.Add(ctx, bytes, attrs) + } +} + +// recordWrite reports one MutateRows RPC: its latency plus the rows and bytes +// it wrote (attempted). It records even on failure, since the request still costs. +func (m *bigtableMetrics) recordWrite(ctx context.Context, table string, latency time.Duration, rows, bytes int64) { + if m == nil { + return + } + attrs := metric.WithAttributes(attribute.String("table", table)) + m.writeLatency.Record(ctx, latency.Seconds(), attrs) + if rows > 0 { + m.rowsMutated.Add(ctx, rows, attrs) + } + if bytes > 0 { + m.bytesWritten.Add(ctx, bytes, attrs) + } +} + +// bigtableRowSize estimates the transferred size of a read row: the row key +// plus each returned cell's qualifier and value. +func bigtableRowSize(row BigtableRow) int64 { + size := int64(len(row.Key)) + for _, cell := range row.Cells { + size += int64(len(cell.Qualifier)) + int64(len(cell.Value)) + } + return size +} + +// bigtableMutationSize estimates the written size of a mutation row: the row +// key plus each set cell's qualifier and value. +func bigtableMutationSize(row BigtableRowMutation) int64 { + size := int64(len(row.RowKey)) + for _, cell := range row.SetCells { + size += int64(len(cell.Qualifier)) + int64(len(cell.Value)) + } + return size +} diff --git a/sei-db/state_db/ss/offload/historical/metrics_test.go b/sei-db/state_db/ss/offload/historical/metrics_test.go new file mode 100644 index 0000000000..d35cfbf608 --- /dev/null +++ b/sei-db/state_db/ss/offload/historical/metrics_test.go @@ -0,0 +1,51 @@ +package historical + +import ( + "context" + "testing" + "time" +) + +func TestBigtableRowSize(t *testing.T) { + row := BigtableRow{ + Key: "abc", // 3 + Cells: []BigtableCell{ + {Qualifier: "value", Value: []byte("hello")}, // 5 + 5 + {Qualifier: "deleted", Value: []byte{0x1}}, // 7 + 1 + }, + } + if got, want := bigtableRowSize(row), int64(3+5+5+7+1); got != want { + t.Fatalf("bigtableRowSize() = %d, want %d", got, want) + } + if got := bigtableRowSize(BigtableRow{}); got != 0 { + t.Fatalf("bigtableRowSize(empty) = %d, want 0", got) + } +} + +func TestBigtableMutationSize(t *testing.T) { + row := BigtableRowMutation{ + RowKey: "row1", // 4 + SetCells: []BigtableSetCell{ + {Qualifier: "value", Value: []byte("data")}, // 5 + 4 + }, + } + if got, want := bigtableMutationSize(row), int64(4+5+4); got != want { + t.Fatalf("bigtableMutationSize() = %d, want %d", got, want) + } +} + +// A nil *bigtableMetrics must be a safe no-op so callers never need nil checks +// and metrics stay zero-overhead when no MeterProvider records them. +func TestBigtableMetricsNilSafe(t *testing.T) { + var m *bigtableMetrics + m.recordRead(context.Background(), "state", time.Millisecond, 2, 128) + m.recordWrite(context.Background(), "state", time.Millisecond, 4, 256) +} + +func TestBigtableMetricsRecordNoPanic(t *testing.T) { + m := newBigtableMetrics() + m.recordRead(context.Background(), "state", 5*time.Millisecond, 3, 512) + m.recordWrite(context.Background(), "state", 7*time.Millisecond, 10, 4096) + // Zero rows/bytes must still record latency without panicking. + m.recordRead(context.Background(), "state", time.Millisecond, 0, 0) +} diff --git a/sei-db/state_db/ss/offload/historical/reader.go b/sei-db/state_db/ss/offload/historical/reader.go new file mode 100644 index 0000000000..45fb25dd10 --- /dev/null +++ b/sei-db/state_db/ss/offload/historical/reader.go @@ -0,0 +1,28 @@ +// Package historical reads MVCC state from an external historical store. +package historical + +import ( + "context" + "errors" +) + +var ErrNotFound = errors.New("historical state not found") + +// Value is the actual MVCC value that satisfied the lookup. +// Version may be older than the requested target version. +type Value struct { + Bytes []byte + Version int64 +} + +type Reader interface { + // Get returns ErrNotFound if no row exists at or before targetVersion, + // or if the latest such row is a tombstone. + Get(ctx context.Context, storeName string, key []byte, targetVersion int64) (Value, error) + + // Has skips value transfer and returns false for missing or tombstoned keys. + Has(ctx context.Context, storeName string, key []byte, targetVersion int64) (bool, error) + + LastVersion(ctx context.Context) (int64, error) + Close() error +} diff --git a/sei-db/state_db/ss/offload/kafka.go b/sei-db/state_db/ss/offload/kafka.go index edbe366818..8bce216c6a 100644 --- a/sei-db/state_db/ss/offload/kafka.go +++ b/sei-db/state_db/ss/offload/kafka.go @@ -12,6 +12,7 @@ import ( "github.com/segmentio/kafka-go" "github.com/segmentio/kafka-go/compress" "github.com/segmentio/kafka-go/sasl" + "github.com/segmentio/kafka-go/sasl/plain" dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" ) @@ -19,18 +20,23 @@ import ( const kafkaOptionNone = "none" type KafkaConfig struct { - Brokers []string - Topic string - ClientID string - Region string - Async bool - RequiredAcks string - Compression string - BatchSize int - BatchTimeout time.Duration - BatchBytes int - TLSEnabled bool + Brokers []string + Topic string + ClientID string + Region string + Async bool + RequiredAcks string + Compression string + BatchSize int + BatchTimeout time.Duration + BatchBytes int + TLSEnabled bool + // SASLMechanism selects broker auth: "none", "plain" (username/password, + // e.g. Google Cloud Managed Kafka service-account credentials), or + // "aws-msk-iam". SASLMechanism string + Username string + Password string } func (c *KafkaConfig) ApplyDefaults() { @@ -83,20 +89,29 @@ func (c *KafkaConfig) Validate() error { return fmt.Errorf("unsupported kafka compression %q", c.Compression) } - switch strings.ToLower(c.SASLMechanism) { + return ValidateSASL(*c) +} + +// ValidateSASL checks the SASL mechanism and the credentials it requires. +// Shared by the producer and the offload consumer configs. +func ValidateSASL(cfg KafkaConfig) error { + switch strings.ToLower(cfg.SASLMechanism) { case "", kafkaOptionNone: - return nil + case "plain": + if cfg.Username == "" || cfg.Password == "" { + return fmt.Errorf("kafka username and password are required for sasl plain") + } case "aws-msk-iam": - if !c.TLSEnabled { + if !cfg.TLSEnabled { return fmt.Errorf("kafka tls must be enabled for aws-msk-iam") } - if c.Region == "" { + if cfg.Region == "" { return fmt.Errorf("kafka region is required for aws-msk-iam") } - return nil default: - return fmt.Errorf("unsupported kafka sasl mechanism %q", c.SASLMechanism) + return fmt.Errorf("unsupported kafka sasl mechanism %q", cfg.SASLMechanism) } + return nil } type kafkaStream struct { @@ -122,7 +137,7 @@ func NewKafkaStream(cfg KafkaConfig) (Stream, error) { } } - mechanism, err := kafkaSASLMechanism(cfg) + mechanism, err := NewSASLMechanism(cfg) if err != nil { return nil, err } @@ -211,10 +226,14 @@ func kafkaCompression(name string) compress.Compression { } } -func kafkaSASLMechanism(cfg KafkaConfig) (sasl.Mechanism, error) { +// NewSASLMechanism returns the SASL mechanism for cfg, which must already +// have passed ValidateSASL. +func NewSASLMechanism(cfg KafkaConfig) (sasl.Mechanism, error) { switch strings.ToLower(cfg.SASLMechanism) { case "", kafkaOptionNone: return nil, nil + case "plain": + return plain.Mechanism{Username: cfg.Username, Password: cfg.Password}, nil case "aws-msk-iam": return newAWSMSKIAMMechanism(cfg) default: From f5ac8bbf2c496961e12b4173a4e9320d7ef78bec Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Tue, 21 Jul 2026 13:44:25 -0400 Subject: [PATCH 2/4] Require all Bigtable connection params in Configured Configured() ORed the three connection fields, so a config with only ProjectID set read as configured and then failed later when the client opened. Require ProjectID, InstanceID, and Table to all be non-blank so partially filled configs are treated as not configured. Co-Authored-By: Claude Fable 5 --- sei-db/state_db/ss/offload/historical/bigtable.go | 5 +++-- sei-db/state_db/ss/offload/historical/bigtable_test.go | 10 ++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/sei-db/state_db/ss/offload/historical/bigtable.go b/sei-db/state_db/ss/offload/historical/bigtable.go index edb85510fa..64b6e4efbc 100644 --- a/sei-db/state_db/ss/offload/historical/bigtable.go +++ b/sei-db/state_db/ss/offload/historical/bigtable.go @@ -115,9 +115,10 @@ func (c *BigtableConfig) ApplyDefaults() { } } +// Configured reports whether all required connection parameters are set. func (c BigtableConfig) Configured() bool { - return strings.TrimSpace(c.ProjectID) != "" || - strings.TrimSpace(c.InstanceID) != "" || + return strings.TrimSpace(c.ProjectID) != "" && + strings.TrimSpace(c.InstanceID) != "" && strings.TrimSpace(c.Table) != "" } diff --git a/sei-db/state_db/ss/offload/historical/bigtable_test.go b/sei-db/state_db/ss/offload/historical/bigtable_test.go index e5a93e9291..fd3b2cee0d 100644 --- a/sei-db/state_db/ss/offload/historical/bigtable_test.go +++ b/sei-db/state_db/ss/offload/historical/bigtable_test.go @@ -35,6 +35,16 @@ func TestBigtableConfigDefaultsAndValidate(t *testing.T) { require.ErrorContains(t, missingShards.Validate(), "shards") } +func TestBigtableConfigConfigured(t *testing.T) { + full := BigtableConfig{ProjectID: "p", InstanceID: "i", Table: "t"} + require.True(t, full.Configured()) + + onlyProject := BigtableConfig{ProjectID: "p"} + require.False(t, onlyProject.Configured()) + + require.False(t, BigtableConfig{}.Configured()) +} + func TestBigtableMutationRowKeyOrdersLatestVersionFirst(t *testing.T) { key40 := BigtableMutationRowKey("bank", []byte("k1"), 40, 256) key60 := BigtableMutationRowKey("bank", []byte("k1"), 60, 256) From 927d9c400ce72bf15ef40b43bd051a13b5bbb01c Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Tue, 21 Jul 2026 13:49:43 -0400 Subject: [PATCH 3/4] Move Kafka transport into sei-db/queue/kafka Pull the Kafka plumbing out of the state-store offload package into a standalone queue transport library: - writer construction (WriterConfig, SASL validation and mechanisms, the AWS MSK IAM signer) moves from offload, - reader construction (ReaderConfig) moves from offload/consumer, - the generic batching consume loop (Consumer, Options, Sink) and its throughput/lag metrics move from offload/consumer. The loop is now payload-agnostic: Sink.WriteBatch receives raw Kafka messages and owns decoding. The Bigtable sink decodes changelog entries itself, so offload/consumer keeps only the changelog-to-Bigtable application. offload.KafkaConfig stays as a type alias so existing callers keep compiling. Meter and metric names are unchanged. Co-Authored-By: Claude Fable 5 --- .../ss/offload => queue/kafka}/aws_msk_iam.go | 4 +- .../consumer => queue/kafka}/consumer.go | 65 +++--- .../consumer => queue/kafka}/metrics.go | 2 +- .../consumer => queue/kafka}/metrics_test.go | 8 +- .../kafka.go => queue/kafka/reader.go} | 42 ++-- .../kafka/reader_test.go} | 30 +-- sei-db/queue/kafka/writer.go | 199 ++++++++++++++++++ sei-db/queue/kafka/writer_test.go | 51 +++++ .../state_db/ss/offload/consumer/bigtable.go | 24 ++- .../ss/offload/consumer/bigtable_test.go | 57 +++-- .../cmd/historical-offload-consumer/main.go | 5 +- sei-db/state_db/ss/offload/consumer/config.go | 4 +- .../ss/offload/consumer/config_test.go | 4 +- sei-db/state_db/ss/offload/consumer/sink.go | 13 +- .../state_db/ss/offload/consumer/sink_test.go | 17 ++ sei-db/state_db/ss/offload/kafka.go | 193 +---------------- sei-db/state_db/ss/offload/kafka_test.go | 43 ---- 17 files changed, 418 insertions(+), 343 deletions(-) rename sei-db/{state_db/ss/offload => queue/kafka}/aws_msk_iam.go (97%) rename sei-db/{state_db/ss/offload/consumer => queue/kafka}/consumer.go (78%) rename sei-db/{state_db/ss/offload/consumer => queue/kafka}/metrics.go (99%) rename sei-db/{state_db/ss/offload/consumer => queue/kafka}/metrics_test.go (86%) rename sei-db/{state_db/ss/offload/consumer/kafka.go => queue/kafka/reader.go} (68%) rename sei-db/{state_db/ss/offload/consumer/kafka_test.go => queue/kafka/reader_test.go} (65%) create mode 100644 sei-db/queue/kafka/writer.go create mode 100644 sei-db/queue/kafka/writer_test.go create mode 100644 sei-db/state_db/ss/offload/consumer/sink_test.go diff --git a/sei-db/state_db/ss/offload/aws_msk_iam.go b/sei-db/queue/kafka/aws_msk_iam.go similarity index 97% rename from sei-db/state_db/ss/offload/aws_msk_iam.go rename to sei-db/queue/kafka/aws_msk_iam.go index ae066d032c..4311a79447 100644 --- a/sei-db/state_db/ss/offload/aws_msk_iam.go +++ b/sei-db/queue/kafka/aws_msk_iam.go @@ -1,4 +1,4 @@ -package offload +package kafka import ( "context" @@ -40,7 +40,7 @@ type awsMSKIAMMechanism struct { var _ sasl.Mechanism = (*awsMSKIAMMechanism)(nil) var _ sasl.StateMachine = (*awsMSKIAMMechanism)(nil) -func newAWSMSKIAMMechanism(cfg KafkaConfig) (sasl.Mechanism, error) { +func newAWSMSKIAMMechanism(cfg WriterConfig) (sasl.Mechanism, error) { sess, err := session.NewSessionWithOptions(session.Options{ Config: aws.Config{ Region: aws.String(cfg.Region), diff --git a/sei-db/state_db/ss/offload/consumer/consumer.go b/sei-db/queue/kafka/consumer.go similarity index 78% rename from sei-db/state_db/ss/offload/consumer/consumer.go rename to sei-db/queue/kafka/consumer.go index d021cd9c6b..5708b667e9 100644 --- a/sei-db/state_db/ss/offload/consumer/consumer.go +++ b/sei-db/queue/kafka/consumer.go @@ -1,4 +1,4 @@ -package consumer +package kafka import ( "context" @@ -6,14 +6,21 @@ import ( "fmt" "time" - "github.com/segmentio/kafka-go" + kafkago "github.com/segmentio/kafka-go" "golang.org/x/sync/errgroup" ) // MessageSource is the subset of *kafka.Reader used by the loop. type MessageSource interface { - FetchMessage(ctx context.Context) (kafka.Message, error) - CommitMessages(ctx context.Context, msgs ...kafka.Message) error + FetchMessage(ctx context.Context) (kafkago.Message, error) + CommitMessages(ctx context.Context, msgs ...kafkago.Message) error +} + +// Sink persists batches of raw Kafka messages. Implementations own payload +// decoding; the consume loop stays payload-agnostic. +type Sink interface { + WriteBatch(ctx context.Context, msgs []kafkago.Message) error + Close() error } // Messages are sharded by partition: cross-partition writes parallelize while @@ -86,9 +93,9 @@ func New(reader MessageSource, sink Sink, opts Options) *Consumer { // Run commits offsets only after the sink persists each message. func (c *Consumer) Run(ctx context.Context) error { g, gctx := errgroup.WithContext(ctx) - shards := make([]chan kafka.Message, c.workers) + shards := make([]chan kafkago.Message, c.workers) for i := range shards { - shards[i] = make(chan kafka.Message, c.shardBuf) + shards[i] = make(chan kafkago.Message, c.shardBuf) ch := shards[i] g.Go(func() error { return c.workerLoop(gctx, ch) }) } @@ -120,7 +127,7 @@ func (c *Consumer) Run(ctx context.Context) error { return nil } -func (c *Consumer) workerLoop(ctx context.Context, ch <-chan kafka.Message) error { +func (c *Consumer) workerLoop(ctx context.Context, ch <-chan kafkago.Message) error { for { select { case <-ctx.Done(): @@ -143,8 +150,8 @@ func (c *Consumer) workerLoop(ctx context.Context, ch <-chan kafka.Message) erro } } -func (c *Consumer) collectBatch(ctx context.Context, ch <-chan kafka.Message, first kafka.Message) ([]kafka.Message, bool) { - msgs := make([]kafka.Message, 1, c.batchSize) +func (c *Consumer) collectBatch(ctx context.Context, ch <-chan kafkago.Message, first kafkago.Message) ([]kafkago.Message, bool) { + msgs := make([]kafkago.Message, 1, c.batchSize) msgs[0] = first if c.batchSize <= 1 { return msgs, true @@ -186,44 +193,28 @@ drainBuffered: return msgs, true } -func (c *Consumer) processBatch(ctx context.Context, msgs []kafka.Message) error { - records := make([]Record, 0, len(msgs)) - var firstVersion, lastVersion int64 - for i, msg := range msgs { - entry, err := DecodeEntry(msg.Value) - if err != nil { - return fmt.Errorf("decode message at offset %d: %w", msg.Offset, err) - } - if i == 0 { - firstVersion = entry.Version - } - lastVersion = entry.Version - records = append(records, Record{ - Topic: msg.Topic, - Partition: msg.Partition, - Offset: msg.Offset, - Entry: entry, - }) - } +func (c *Consumer) processBatch(ctx context.Context, msgs []kafkago.Message) error { + firstOffset := msgs[0].Offset + lastOffset := msgs[len(msgs)-1].Offset start := time.Now() - if err := c.writeBatchWithRetry(ctx, records); err != nil { - return fmt.Errorf("sink write batch first_version=%d last_version=%d: %w", - firstVersion, lastVersion, err) + if err := c.writeBatchWithRetry(ctx, msgs); err != nil { + return fmt.Errorf("sink write batch first_offset=%d last_offset=%d: %w", + firstOffset, lastOffset, err) } writeLatency := time.Since(start) - c.logf("wrote records=%d first_version=%d last_version=%d in %s", - len(records), firstVersion, lastVersion, writeLatency) + c.logf("wrote messages=%d first_offset=%d last_offset=%d in %s", + len(msgs), firstOffset, lastOffset, writeLatency) if err := c.reader.CommitMessages(ctx, msgs...); err != nil { return fmt.Errorf("commit kafka offsets: %w", err) } - c.metrics.recordBatch(ctx, int64(len(records)), batchLag(msgs), writeLatency) + c.metrics.recordBatch(ctx, int64(len(msgs)), batchLag(msgs), writeLatency) return nil } // batchLag reports how far behind the partition head the batch left the // consumer: the max over messages of (high watermark - offset - 1). Returns -1 // when no message carries a watermark so the gauge is left untouched. -func batchLag(msgs []kafka.Message) int64 { +func batchLag(msgs []kafkago.Message) int64 { lag := int64(-1) for _, msg := range msgs { if msg.HighWaterMark <= 0 { @@ -236,11 +227,11 @@ func batchLag(msgs []kafka.Message) int64 { return lag } -func (c *Consumer) writeBatchWithRetry(ctx context.Context, records []Record) error { +func (c *Consumer) writeBatchWithRetry(ctx context.Context, msgs []kafkago.Message) error { backoff := sinkBaseBackoff var lastErr error for attempt := 1; attempt <= sinkMaxAttempts; attempt++ { - err := c.sink.WriteBatch(ctx, records) + err := c.sink.WriteBatch(ctx, msgs) if err == nil { return nil } diff --git a/sei-db/state_db/ss/offload/consumer/metrics.go b/sei-db/queue/kafka/metrics.go similarity index 99% rename from sei-db/state_db/ss/offload/consumer/metrics.go rename to sei-db/queue/kafka/metrics.go index afcca00702..8cafc2c28a 100644 --- a/sei-db/state_db/ss/offload/consumer/metrics.go +++ b/sei-db/queue/kafka/metrics.go @@ -1,4 +1,4 @@ -package consumer +package kafka import ( "context" diff --git a/sei-db/state_db/ss/offload/consumer/metrics_test.go b/sei-db/queue/kafka/metrics_test.go similarity index 86% rename from sei-db/state_db/ss/offload/consumer/metrics_test.go rename to sei-db/queue/kafka/metrics_test.go index 7b77d3b4dc..59fd39de07 100644 --- a/sei-db/state_db/ss/offload/consumer/metrics_test.go +++ b/sei-db/queue/kafka/metrics_test.go @@ -1,18 +1,18 @@ -package consumer +package kafka import ( "context" "testing" "time" - "github.com/segmentio/kafka-go" + kafkago "github.com/segmentio/kafka-go" "github.com/stretchr/testify/require" ) func TestBatchLag(t *testing.T) { // Max over the batch of (HighWaterMark - Offset - 1); the head message // (offset == HWM-1) has lag 0. - msgs := []kafka.Message{ + msgs := []kafkago.Message{ {Offset: 10, HighWaterMark: 100}, // 89 {Offset: 98, HighWaterMark: 100}, // 1 {Offset: 99, HighWaterMark: 100}, // 0 (head) @@ -20,7 +20,7 @@ func TestBatchLag(t *testing.T) { require.Equal(t, int64(89), batchLag(msgs)) // No watermark information -> sentinel -1 so the gauge is left untouched. - require.Equal(t, int64(-1), batchLag([]kafka.Message{{Offset: 5}})) + require.Equal(t, int64(-1), batchLag([]kafkago.Message{{Offset: 5}})) require.Equal(t, int64(-1), batchLag(nil)) } diff --git a/sei-db/state_db/ss/offload/consumer/kafka.go b/sei-db/queue/kafka/reader.go similarity index 68% rename from sei-db/state_db/ss/offload/consumer/kafka.go rename to sei-db/queue/kafka/reader.go index 894fbea0ce..e20835992a 100644 --- a/sei-db/state_db/ss/offload/consumer/kafka.go +++ b/sei-db/queue/kafka/reader.go @@ -1,4 +1,4 @@ -package consumer +package kafka import ( "crypto/tls" @@ -6,17 +6,13 @@ import ( "strings" "time" - gogoproto "github.com/gogo/protobuf/proto" - "github.com/segmentio/kafka-go" - - dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" - "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload" + kafkago "github.com/segmentio/kafka-go" ) // TLS/SASL must match the producer cluster. Commits are synchronous // (kafka-go's zero CommitInterval) so offsets only advance after the sink // persists each entry. -type KafkaReaderConfig struct { +type ReaderConfig struct { Brokers []string Topic string GroupID string @@ -35,7 +31,7 @@ type KafkaReaderConfig struct { Password string } -func (c *KafkaReaderConfig) ApplyDefaults() { +func (c *ReaderConfig) ApplyDefaults() { if c.ClientID == "" { c.ClientID = "sei-historical-offload-consumer" } @@ -53,7 +49,7 @@ func (c *KafkaReaderConfig) ApplyDefaults() { } } -func (c *KafkaReaderConfig) Validate() error { +func (c *ReaderConfig) Validate() error { if len(c.Brokers) == 0 { return fmt.Errorf("kafka brokers are required") } @@ -68,11 +64,11 @@ func (c *KafkaReaderConfig) Validate() error { default: return fmt.Errorf("unsupported kafka start offset %q", c.StartOffset) } - return offload.ValidateSASL(c.saslConfig()) + return ValidateSASL(c.saslConfig()) } -func (c KafkaReaderConfig) saslConfig() offload.KafkaConfig { - return offload.KafkaConfig{ +func (c ReaderConfig) saslConfig() WriterConfig { + return WriterConfig{ Region: c.Region, TLSEnabled: c.TLSEnabled, SASLMechanism: c.SASLMechanism, @@ -81,31 +77,33 @@ func (c KafkaReaderConfig) saslConfig() offload.KafkaConfig { } } -func NewKafkaReader(cfg KafkaReaderConfig) (*kafka.Reader, error) { +// NewReader builds a kafka-go Reader from cfg after applying defaults and +// validating it. +func NewReader(cfg ReaderConfig) (*kafkago.Reader, error) { cfg.ApplyDefaults() if err := cfg.Validate(); err != nil { return nil, err } - dialer := &kafka.Dialer{ + dialer := &kafkago.Dialer{ ClientID: cfg.ClientID, Timeout: 10 * time.Second, } if cfg.TLSEnabled { dialer.TLS = &tls.Config{MinVersion: tls.VersionTLS12} } - mech, err := offload.NewSASLMechanism(cfg.saslConfig()) + mech, err := NewSASLMechanism(cfg.saslConfig()) if err != nil { return nil, err } dialer.SASLMechanism = mech - start := kafka.FirstOffset + start := kafkago.FirstOffset if strings.EqualFold(cfg.StartOffset, "last") { - start = kafka.LastOffset + start = kafkago.LastOffset } - return kafka.NewReader(kafka.ReaderConfig{ + return kafkago.NewReader(kafkago.ReaderConfig{ Brokers: cfg.Brokers, Topic: cfg.Topic, GroupID: cfg.GroupID, @@ -116,11 +114,3 @@ func NewKafkaReader(cfg KafkaReaderConfig) (*kafka.Reader, error) { StartOffset: start, }), nil } - -func DecodeEntry(payload []byte) (*dbproto.ChangelogEntry, error) { - entry := &dbproto.ChangelogEntry{} - if err := gogoproto.Unmarshal(payload, entry); err != nil { - return nil, fmt.Errorf("decode changelog entry: %w", err) - } - return entry, nil -} diff --git a/sei-db/state_db/ss/offload/consumer/kafka_test.go b/sei-db/queue/kafka/reader_test.go similarity index 65% rename from sei-db/state_db/ss/offload/consumer/kafka_test.go rename to sei-db/queue/kafka/reader_test.go index 57056e9951..a393bc2cfd 100644 --- a/sei-db/state_db/ss/offload/consumer/kafka_test.go +++ b/sei-db/queue/kafka/reader_test.go @@ -1,14 +1,13 @@ -package consumer +package kafka import ( "testing" - "github.com/sei-protocol/sei-chain/sei-db/proto" "github.com/stretchr/testify/require" ) -func TestKafkaReaderConfigApplyDefaults(t *testing.T) { - cfg := KafkaReaderConfig{ +func TestReaderConfigApplyDefaults(t *testing.T) { + cfg := ReaderConfig{ Brokers: []string{"localhost:9092"}, Topic: "historical-offload", GroupID: "bigtable", @@ -20,14 +19,14 @@ func TestKafkaReaderConfigApplyDefaults(t *testing.T) { require.Equal(t, 10<<20, cfg.MaxBytes) } -func TestKafkaReaderConfigValidate(t *testing.T) { - cfg := KafkaReaderConfig{} +func TestReaderConfigValidate(t *testing.T) { + cfg := ReaderConfig{} require.ErrorContains(t, cfg.Validate(), "brokers") - cfg = KafkaReaderConfig{Brokers: []string{"x"}} + cfg = ReaderConfig{Brokers: []string{"x"}} require.ErrorContains(t, cfg.Validate(), "topic") - cfg = KafkaReaderConfig{Brokers: []string{"x"}, Topic: "t"} + cfg = ReaderConfig{Brokers: []string{"x"}, Topic: "t"} require.ErrorContains(t, cfg.Validate(), "group id") - cfg = KafkaReaderConfig{ + cfg = ReaderConfig{ Brokers: []string{"x"}, Topic: "t", GroupID: "g", @@ -36,8 +35,8 @@ func TestKafkaReaderConfigValidate(t *testing.T) { require.ErrorContains(t, cfg.Validate(), "start offset") } -func TestKafkaReaderConfigValidateSASL(t *testing.T) { - base := KafkaReaderConfig{Brokers: []string{"x"}, Topic: "t", GroupID: "g"} +func TestReaderConfigValidateSASL(t *testing.T) { + base := ReaderConfig{Brokers: []string{"x"}, Topic: "t", GroupID: "g"} cfg := base cfg.SASLMechanism = "aws-msk-iam" @@ -66,12 +65,3 @@ func TestKafkaReaderConfigValidateSASL(t *testing.T) { cfg.SASLMechanism = "none" require.NoError(t, cfg.Validate()) } - -func TestDecodeEntry(t *testing.T) { - entry := &proto.ChangelogEntry{Version: 7} - payload, err := entry.Marshal() - require.NoError(t, err) - got, err := DecodeEntry(payload) - require.NoError(t, err) - require.Equal(t, int64(7), got.Version) -} diff --git a/sei-db/queue/kafka/writer.go b/sei-db/queue/kafka/writer.go new file mode 100644 index 0000000000..2f360fc98c --- /dev/null +++ b/sei-db/queue/kafka/writer.go @@ -0,0 +1,199 @@ +// Package kafka bundles the Kafka transport used by the historical state +// offload: producer (writer) construction, consumer (reader) construction, +// and the generic batching consume loop that drains a topic into a Sink. +package kafka + +import ( + "crypto/tls" + "fmt" + "strings" + "time" + + kafkago "github.com/segmentio/kafka-go" + "github.com/segmentio/kafka-go/compress" + "github.com/segmentio/kafka-go/sasl" + "github.com/segmentio/kafka-go/sasl/plain" +) + +const kafkaOptionNone = "none" + +type WriterConfig struct { + Brokers []string + Topic string + ClientID string + Region string + Async bool + RequiredAcks string + Compression string + BatchSize int + BatchTimeout time.Duration + BatchBytes int + TLSEnabled bool + // SASLMechanism selects broker auth: "none", "plain" (username/password, + // e.g. Google Cloud Managed Kafka service-account credentials), or + // "aws-msk-iam". + SASLMechanism string + Username string + Password string +} + +func (c *WriterConfig) ApplyDefaults() { + if c.ClientID == "" { + c.ClientID = "cryptosim-historical-offload" + } + if c.RequiredAcks == "" { + c.RequiredAcks = kafkaOptionNone + } + if c.Compression == "" { + c.Compression = "snappy" + } + if c.BatchSize == 0 { + c.BatchSize = 1000 + } + if c.BatchTimeout == 0 { + c.BatchTimeout = 50 * time.Millisecond + } + if c.BatchBytes == 0 { + c.BatchBytes = 4 << 20 + } +} + +func (c *WriterConfig) Validate() error { + if len(c.Brokers) == 0 { + return fmt.Errorf("kafka brokers are required") + } + if c.Topic == "" { + return fmt.Errorf("kafka topic is required") + } + if c.BatchTimeout < 0 { + return fmt.Errorf("kafka batch timeout must be non-negative") + } + if c.BatchSize < 0 { + return fmt.Errorf("kafka batch size must be non-negative") + } + if c.BatchBytes < 0 { + return fmt.Errorf("kafka batch bytes must be non-negative") + } + + switch strings.ToLower(c.RequiredAcks) { + case kafkaOptionNone, "leader", "all": + default: + return fmt.Errorf("unsupported kafka required acks %q", c.RequiredAcks) + } + + switch strings.ToLower(c.Compression) { + case kafkaOptionNone, "gzip", "snappy", "lz4", "zstd": + default: + return fmt.Errorf("unsupported kafka compression %q", c.Compression) + } + + return ValidateSASL(*c) +} + +// ValidateSASL checks the SASL mechanism and the credentials it requires. +// Shared by the producer and the offload consumer configs. +func ValidateSASL(cfg WriterConfig) error { + switch strings.ToLower(cfg.SASLMechanism) { + case "", kafkaOptionNone: + case "plain": + if cfg.Username == "" || cfg.Password == "" { + return fmt.Errorf("kafka username and password are required for sasl plain") + } + case "aws-msk-iam": + if !cfg.TLSEnabled { + return fmt.Errorf("kafka tls must be enabled for aws-msk-iam") + } + if cfg.Region == "" { + return fmt.Errorf("kafka region is required for aws-msk-iam") + } + default: + return fmt.Errorf("unsupported kafka sasl mechanism %q", cfg.SASLMechanism) + } + return nil +} + +// NewWriter builds a kafka-go Writer from cfg after applying defaults and +// validating it. +func NewWriter(cfg WriterConfig) (*kafkago.Writer, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + + dialer := &kafkago.Dialer{ + ClientID: cfg.ClientID, + Timeout: 10 * time.Second, + } + if cfg.TLSEnabled { + dialer.TLS = &tls.Config{ + MinVersion: tls.VersionTLS12, + } + } + + mechanism, err := NewSASLMechanism(cfg) + if err != nil { + return nil, err + } + dialer.SASLMechanism = mechanism + + return &kafkago.Writer{ + Addr: kafkago.TCP(cfg.Brokers...), + Topic: cfg.Topic, + Balancer: &kafkago.Hash{}, + RequiredAcks: kafkaRequiredAcks(cfg.RequiredAcks), + BatchSize: cfg.BatchSize, + BatchTimeout: cfg.BatchTimeout, + BatchBytes: int64(cfg.BatchBytes), + Async: cfg.Async, + Compression: kafkaCompression(cfg.Compression), + Transport: &kafkago.Transport{ + ClientID: cfg.ClientID, + TLS: dialer.TLS, + SASL: mechanism, + IdleTimeout: 30 * time.Second, + MetadataTTL: 30 * time.Second, + DialTimeout: dialer.Timeout, + }, + }, nil +} + +func kafkaRequiredAcks(requiredAcks string) kafkago.RequiredAcks { + switch strings.ToLower(requiredAcks) { + case kafkaOptionNone: + return kafkago.RequireNone + case "leader": + return kafkago.RequireOne + default: + return kafkago.RequireAll + } +} + +func kafkaCompression(name string) compress.Compression { + switch strings.ToLower(name) { + case "gzip": + return compress.Gzip + case "lz4": + return compress.Lz4 + case "zstd": + return compress.Zstd + case kafkaOptionNone: + return compress.None + default: + return compress.Snappy + } +} + +// NewSASLMechanism returns the SASL mechanism for cfg, which must already +// have passed ValidateSASL. +func NewSASLMechanism(cfg WriterConfig) (sasl.Mechanism, error) { + switch strings.ToLower(cfg.SASLMechanism) { + case "", kafkaOptionNone: + return nil, nil + case "plain": + return plain.Mechanism{Username: cfg.Username, Password: cfg.Password}, nil + case "aws-msk-iam": + return newAWSMSKIAMMechanism(cfg) + default: + return nil, fmt.Errorf("unsupported kafka sasl mechanism %q", cfg.SASLMechanism) + } +} diff --git a/sei-db/queue/kafka/writer_test.go b/sei-db/queue/kafka/writer_test.go new file mode 100644 index 0000000000..f789a94d28 --- /dev/null +++ b/sei-db/queue/kafka/writer_test.go @@ -0,0 +1,51 @@ +package kafka + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestWriterConfigApplyDefaultsAndValidate(t *testing.T) { + cfg := WriterConfig{ + Brokers: []string{"localhost:9092"}, + Topic: "historical-offload", + } + + cfg.ApplyDefaults() + require.Equal(t, "cryptosim-historical-offload", cfg.ClientID) + require.Equal(t, "none", cfg.RequiredAcks) + require.Equal(t, "snappy", cfg.Compression) + require.Equal(t, 1000, cfg.BatchSize) + require.Equal(t, 4<<20, cfg.BatchBytes) + require.NoError(t, cfg.Validate()) +} + +func TestWriterConfigValidateRequiresBrokerAndTopic(t *testing.T) { + cfg := WriterConfig{} + require.Error(t, cfg.Validate()) +} + +func TestWriterConfigValidateRequiresRegionAndTLSForAWSMSKIAM(t *testing.T) { + cfg := WriterConfig{ + Brokers: []string{"localhost:9098"}, + Topic: "historical-offload", + SASLMechanism: "aws-msk-iam", + } + cfg.ApplyDefaults() + + require.ErrorContains(t, cfg.Validate(), "tls must be enabled") + + cfg.TLSEnabled = true + require.ErrorContains(t, cfg.Validate(), "region is required") + + cfg.Region = "eu-central-1" + require.NoError(t, cfg.Validate()) +} + +func TestAWSMSKIAMMechanismRequiresMetadata(t *testing.T) { + mech := &awsMSKIAMMechanism{region: "eu-central-1"} + _, _, err := mech.Start(context.Background()) + require.ErrorContains(t, err, "missing sasl metadata") +} diff --git a/sei-db/state_db/ss/offload/consumer/bigtable.go b/sei-db/state_db/ss/offload/consumer/bigtable.go index 0e70ce46a3..c362d21d4e 100644 --- a/sei-db/state_db/ss/offload/consumer/bigtable.go +++ b/sei-db/state_db/ss/offload/consumer/bigtable.go @@ -7,9 +7,12 @@ import ( "strconv" "time" + kafkago "github.com/segmentio/kafka-go" + "golang.org/x/sync/errgroup" + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" - "golang.org/x/sync/errgroup" ) type BigtableConfig = historical.BigtableConfig @@ -28,9 +31,9 @@ type bigtableSink struct { bulkChunkWorkers int } -var _ Sink = (*bigtableSink)(nil) +var _ kafka.Sink = (*bigtableSink)(nil) -func NewBigtableSink(cfg BigtableConfig) (Sink, error) { +func NewBigtableSink(cfg BigtableConfig) (kafka.Sink, error) { cfg.ApplyDefaults() if err := cfg.Validate(); err != nil { return nil, err @@ -57,7 +60,20 @@ func (s *bigtableSink) Close() error { return nil } -func (s *bigtableSink) WriteBatch(ctx context.Context, records []Record) error { +func (s *bigtableSink) WriteBatch(ctx context.Context, msgs []kafkago.Message) error { + records := make([]Record, 0, len(msgs)) + for _, msg := range msgs { + entry, err := DecodeEntry(msg.Value) + if err != nil { + return fmt.Errorf("decode message at offset %d: %w", msg.Offset, err) + } + records = append(records, Record{ + Topic: msg.Topic, + Partition: msg.Partition, + Offset: msg.Offset, + Entry: entry, + }) + } records = compactRecords(records) if len(records) == 0 { return nil diff --git a/sei-db/state_db/ss/offload/consumer/bigtable_test.go b/sei-db/state_db/ss/offload/consumer/bigtable_test.go index 63bae0d7b2..28706ef822 100644 --- a/sei-db/state_db/ss/offload/consumer/bigtable_test.go +++ b/sei-db/state_db/ss/offload/consumer/bigtable_test.go @@ -6,11 +6,21 @@ import ( "testing" "time" + kafkago "github.com/segmentio/kafka-go" + "github.com/stretchr/testify/require" + "github.com/sei-protocol/sei-chain/sei-db/proto" "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" - "github.com/stretchr/testify/require" ) +// entryMessage marshals entry into a Kafka message the sink can decode. +func entryMessage(t *testing.T, entry *proto.ChangelogEntry) kafkago.Message { + t.Helper() + payload, err := entry.Marshal() + require.NoError(t, err) + return kafkago.Message{Value: payload} +} + func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { var rows []string sink := &bigtableSink{ @@ -37,7 +47,9 @@ func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { Upgrades: []*proto.TreeNameUpgrade{{Name: "new-store"}}, } - require.NoError(t, sink.WriteBatch(context.Background(), []Record{{Topic: "t", Partition: 1, Offset: 2, Entry: entry}})) + msg := entryMessage(t, entry) + msg.Topic, msg.Partition, msg.Offset = "t", 1, 2 + require.NoError(t, sink.WriteBatch(context.Background(), []kafkago.Message{msg})) require.Len(t, rows, 4) require.ElementsMatch(t, []string{ historical.BigtableMutationRowKey("bank", []byte("k1"), 7, historical.DefaultBigtableShards), @@ -47,6 +59,23 @@ func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { require.Equal(t, historical.BigtableVersionRowKey(7), rows[3]) } +func TestBigtableSinkWriteBatchRejectsUndecodableMessage(t *testing.T) { + applyBulkCalls := 0 + sink := &bigtableSink{ + family: historical.DefaultBigtableFamily, + shards: historical.DefaultBigtableShards, + bulkChunkWorkers: 1, + applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + applyBulkCalls++ + return make([]error, len(mutations)), nil + }, + } + + err := sink.WriteBatch(context.Background(), []kafkago.Message{{Offset: 42, Value: []byte{0xff, 0xff}}}) + require.ErrorContains(t, err, "decode message at offset 42") + require.Zero(t, applyBulkCalls) +} + func TestBigtableSinkWriteBatchWritesRowsBeforeMarkers(t *testing.T) { var calls [][]string @@ -63,24 +92,24 @@ func TestBigtableSinkWriteBatchWritesRowsBeforeMarkers(t *testing.T) { return make([]error, len(mutations)), nil }, } - records := []Record{ - {Entry: &proto.ChangelogEntry{ + msgs := []kafkago.Message{ + entryMessage(t, &proto.ChangelogEntry{ Version: 1, Changesets: []*proto.NamedChangeSet{{ Name: "bank", Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{{Key: []byte("k1"), Value: []byte("v1")}}}, }}, - }}, - {Entry: &proto.ChangelogEntry{ + }), + entryMessage(t, &proto.ChangelogEntry{ Version: 2, Changesets: []*proto.NamedChangeSet{{ Name: "bank", Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{{Key: []byte("k2"), Value: []byte("v2")}}}, }}, - }}, + }), } - require.NoError(t, sink.WriteBatch(context.Background(), records)) + require.NoError(t, sink.WriteBatch(context.Background(), msgs)) require.GreaterOrEqual(t, len(calls), 2) require.ElementsMatch(t, []string{ historical.BigtableMutationRowKey("bank", []byte("k1"), 1, historical.DefaultBigtableShards), @@ -138,8 +167,8 @@ func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { return make([]error, len(mutations)), nil }, } - records := []Record{ - {Entry: &proto.ChangelogEntry{ + msgs := []kafkago.Message{ + entryMessage(t, &proto.ChangelogEntry{ Version: 1, Changesets: []*proto.NamedChangeSet{{ Name: "bank", @@ -149,8 +178,8 @@ func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { }}, }}, Upgrades: []*proto.TreeNameUpgrade{{Name: "new-store"}}, - }}, - {Entry: &proto.ChangelogEntry{ + }), + entryMessage(t, &proto.ChangelogEntry{ Version: 2, Changesets: []*proto.NamedChangeSet{{ Name: "bank", @@ -158,10 +187,10 @@ func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { {Key: []byte("k3"), Value: []byte("v3")}, }}, }}, - }}, + }), } - require.NoError(t, sink.WriteBatch(context.Background(), records)) + require.NoError(t, sink.WriteBatch(context.Background(), msgs)) require.Greater(t, len(calls), 2) markerCall := calls[len(calls)-1] diff --git a/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go b/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go index 1fc43636ce..b8691c46b2 100644 --- a/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go +++ b/sei-db/state_db/ss/offload/consumer/cmd/historical-offload-consumer/main.go @@ -10,6 +10,7 @@ import ( "time" commonmetrics "github.com/sei-protocol/sei-chain/sei-db/common/metrics" + "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/consumer" ) @@ -49,13 +50,13 @@ func main() { } defer func() { _ = sink.Close() }() - reader, err := consumer.NewKafkaReader(cfg.Kafka) + reader, err := kafka.NewReader(cfg.Kafka) if err != nil { log.Fatalf("open kafka reader: %v", err) } defer func() { _ = reader.Close() }() - c := consumer.New(reader, sink, consumer.Options{ + c := kafka.New(reader, sink, kafka.Options{ Logf: func(format string, args ...interface{}) { log.Printf(format, args...) }, Workers: cfg.Workers, ShardBufferSize: cfg.ShardBufferSize, diff --git a/sei-db/state_db/ss/offload/consumer/config.go b/sei-db/state_db/ss/offload/consumer/config.go index 18889f6bf1..43df864425 100644 --- a/sei-db/state_db/ss/offload/consumer/config.go +++ b/sei-db/state_db/ss/offload/consumer/config.go @@ -4,6 +4,8 @@ import ( "encoding/json" "fmt" "os" + + "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" ) const ( @@ -12,7 +14,7 @@ const ( ) type Config struct { - Kafka KafkaReaderConfig + Kafka kafka.ReaderConfig Bigtable BigtableConfig Workers int ShardBufferSize int diff --git a/sei-db/state_db/ss/offload/consumer/config_test.go b/sei-db/state_db/ss/offload/consumer/config_test.go index 9cd6f65d3f..96bc34d909 100644 --- a/sei-db/state_db/ss/offload/consumer/config_test.go +++ b/sei-db/state_db/ss/offload/consumer/config_test.go @@ -6,11 +6,13 @@ import ( "testing" "github.com/stretchr/testify/require" + + "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" ) func TestConfigValidate(t *testing.T) { cfg := Config{ - Kafka: KafkaReaderConfig{ + Kafka: kafka.ReaderConfig{ Brokers: []string{"localhost:9092"}, Topic: "historical-offload", GroupID: "historical-bigtable", diff --git a/sei-db/state_db/ss/offload/consumer/sink.go b/sei-db/state_db/ss/offload/consumer/sink.go index deef46c336..b3d7db76b5 100644 --- a/sei-db/state_db/ss/offload/consumer/sink.go +++ b/sei-db/state_db/ss/offload/consumer/sink.go @@ -1,7 +1,9 @@ package consumer import ( - "context" + "fmt" + + gogoproto "github.com/gogo/protobuf/proto" dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" ) @@ -15,7 +17,10 @@ type Record struct { Entry *dbproto.ChangelogEntry } -type Sink interface { - WriteBatch(ctx context.Context, records []Record) error - Close() error +func DecodeEntry(payload []byte) (*dbproto.ChangelogEntry, error) { + entry := &dbproto.ChangelogEntry{} + if err := gogoproto.Unmarshal(payload, entry); err != nil { + return nil, fmt.Errorf("decode changelog entry: %w", err) + } + return entry, nil } diff --git a/sei-db/state_db/ss/offload/consumer/sink_test.go b/sei-db/state_db/ss/offload/consumer/sink_test.go new file mode 100644 index 0000000000..bc15e913ee --- /dev/null +++ b/sei-db/state_db/ss/offload/consumer/sink_test.go @@ -0,0 +1,17 @@ +package consumer + +import ( + "testing" + + "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/stretchr/testify/require" +) + +func TestDecodeEntry(t *testing.T) { + entry := &proto.ChangelogEntry{Version: 7} + payload, err := entry.Marshal() + require.NoError(t, err) + got, err := DecodeEntry(payload) + require.NoError(t, err) + require.Equal(t, int64(7), got.Version) +} diff --git a/sei-db/state_db/ss/offload/kafka.go b/sei-db/state_db/ss/offload/kafka.go index 8bce216c6a..066402f720 100644 --- a/sei-db/state_db/ss/offload/kafka.go +++ b/sei-db/state_db/ss/offload/kafka.go @@ -2,120 +2,23 @@ package offload import ( "context" - "crypto/tls" "fmt" "strconv" - "strings" "time" gogoproto "github.com/gogo/protobuf/proto" - "github.com/segmentio/kafka-go" - "github.com/segmentio/kafka-go/compress" - "github.com/segmentio/kafka-go/sasl" - "github.com/segmentio/kafka-go/sasl/plain" + kafkago "github.com/segmentio/kafka-go" dbproto "github.com/sei-protocol/sei-chain/sei-db/proto" + "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" ) -const kafkaOptionNone = "none" - -type KafkaConfig struct { - Brokers []string - Topic string - ClientID string - Region string - Async bool - RequiredAcks string - Compression string - BatchSize int - BatchTimeout time.Duration - BatchBytes int - TLSEnabled bool - // SASLMechanism selects broker auth: "none", "plain" (username/password, - // e.g. Google Cloud Managed Kafka service-account credentials), or - // "aws-msk-iam". - SASLMechanism string - Username string - Password string -} - -func (c *KafkaConfig) ApplyDefaults() { - if c.ClientID == "" { - c.ClientID = "cryptosim-historical-offload" - } - if c.RequiredAcks == "" { - c.RequiredAcks = kafkaOptionNone - } - if c.Compression == "" { - c.Compression = "snappy" - } - if c.BatchSize == 0 { - c.BatchSize = 1000 - } - if c.BatchTimeout == 0 { - c.BatchTimeout = 50 * time.Millisecond - } - if c.BatchBytes == 0 { - c.BatchBytes = 4 << 20 - } -} - -func (c *KafkaConfig) Validate() error { - if len(c.Brokers) == 0 { - return fmt.Errorf("kafka brokers are required") - } - if c.Topic == "" { - return fmt.Errorf("kafka topic is required") - } - if c.BatchTimeout < 0 { - return fmt.Errorf("kafka batch timeout must be non-negative") - } - if c.BatchSize < 0 { - return fmt.Errorf("kafka batch size must be non-negative") - } - if c.BatchBytes < 0 { - return fmt.Errorf("kafka batch bytes must be non-negative") - } - - switch strings.ToLower(c.RequiredAcks) { - case kafkaOptionNone, "leader", "all": - default: - return fmt.Errorf("unsupported kafka required acks %q", c.RequiredAcks) - } - - switch strings.ToLower(c.Compression) { - case kafkaOptionNone, "gzip", "snappy", "lz4", "zstd": - default: - return fmt.Errorf("unsupported kafka compression %q", c.Compression) - } - - return ValidateSASL(*c) -} - -// ValidateSASL checks the SASL mechanism and the credentials it requires. -// Shared by the producer and the offload consumer configs. -func ValidateSASL(cfg KafkaConfig) error { - switch strings.ToLower(cfg.SASLMechanism) { - case "", kafkaOptionNone: - case "plain": - if cfg.Username == "" || cfg.Password == "" { - return fmt.Errorf("kafka username and password are required for sasl plain") - } - case "aws-msk-iam": - if !cfg.TLSEnabled { - return fmt.Errorf("kafka tls must be enabled for aws-msk-iam") - } - if cfg.Region == "" { - return fmt.Errorf("kafka region is required for aws-msk-iam") - } - default: - return fmt.Errorf("unsupported kafka sasl mechanism %q", cfg.SASLMechanism) - } - return nil -} +// KafkaConfig aliases the shared Kafka writer configuration so existing +// callers keep building it with keyed literals. +type KafkaConfig = kafka.WriterConfig type kafkaStream struct { - writer *kafka.Writer + writer *kafkago.Writer durable bool } @@ -123,50 +26,13 @@ var _ Stream = (*kafkaStream)(nil) func NewKafkaStream(cfg KafkaConfig) (Stream, error) { cfg.ApplyDefaults() - if err := cfg.Validate(); err != nil { - return nil, err - } - - dialer := &kafka.Dialer{ - ClientID: cfg.ClientID, - Timeout: 10 * time.Second, - } - if cfg.TLSEnabled { - dialer.TLS = &tls.Config{ - MinVersion: tls.VersionTLS12, - } - } - - mechanism, err := NewSASLMechanism(cfg) + writer, err := kafka.NewWriter(cfg) if err != nil { return nil, err } - dialer.SASLMechanism = mechanism - - requiredAcks := kafkaRequiredAcks(cfg.RequiredAcks) - writer := &kafka.Writer{ - Addr: kafka.TCP(cfg.Brokers...), - Topic: cfg.Topic, - Balancer: &kafka.Hash{}, - RequiredAcks: requiredAcks, - BatchSize: cfg.BatchSize, - BatchTimeout: cfg.BatchTimeout, - BatchBytes: int64(cfg.BatchBytes), - Async: cfg.Async, - Compression: kafkaCompression(cfg.Compression), - Transport: &kafka.Transport{ - ClientID: cfg.ClientID, - TLS: dialer.TLS, - SASL: mechanism, - IdleTimeout: 30 * time.Second, - MetadataTTL: 30 * time.Second, - DialTimeout: dialer.Timeout, - }, - } - return &kafkaStream{ writer: writer, - durable: !cfg.Async && requiredAcks == kafka.RequireAll, + durable: !cfg.Async && writer.RequiredAcks == kafkago.RequireAll, }, nil } @@ -181,7 +47,7 @@ func (k *kafkaStream) Publish(ctx context.Context, entry *dbproto.ChangelogEntry } key := []byte(strconv.FormatInt(entry.Version, 10)) - if err := k.writer.WriteMessages(ctx, kafka.Message{ + if err := k.writer.WriteMessages(ctx, kafkago.Message{ Key: key, Value: payload, Time: time.Now(), @@ -199,44 +65,3 @@ func (k *kafkaStream) Publish(ctx context.Context, entry *dbproto.ChangelogEntry func (k *kafkaStream) Close() error { return k.writer.Close() } - -func kafkaRequiredAcks(requiredAcks string) kafka.RequiredAcks { - switch strings.ToLower(requiredAcks) { - case kafkaOptionNone: - return kafka.RequireNone - case "leader": - return kafka.RequireOne - default: - return kafka.RequireAll - } -} - -func kafkaCompression(name string) compress.Compression { - switch strings.ToLower(name) { - case "gzip": - return compress.Gzip - case "lz4": - return compress.Lz4 - case "zstd": - return compress.Zstd - case kafkaOptionNone: - return compress.None - default: - return compress.Snappy - } -} - -// NewSASLMechanism returns the SASL mechanism for cfg, which must already -// have passed ValidateSASL. -func NewSASLMechanism(cfg KafkaConfig) (sasl.Mechanism, error) { - switch strings.ToLower(cfg.SASLMechanism) { - case "", kafkaOptionNone: - return nil, nil - case "plain": - return plain.Mechanism{Username: cfg.Username, Password: cfg.Password}, nil - case "aws-msk-iam": - return newAWSMSKIAMMechanism(cfg) - default: - return nil, fmt.Errorf("unsupported kafka sasl mechanism %q", cfg.SASLMechanism) - } -} diff --git a/sei-db/state_db/ss/offload/kafka_test.go b/sei-db/state_db/ss/offload/kafka_test.go index 53f3539dbc..2af41c456d 100644 --- a/sei-db/state_db/ss/offload/kafka_test.go +++ b/sei-db/state_db/ss/offload/kafka_test.go @@ -7,43 +7,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestKafkaConfigApplyDefaultsAndValidate(t *testing.T) { - cfg := KafkaConfig{ - Brokers: []string{"localhost:9092"}, - Topic: "historical-offload", - } - - cfg.ApplyDefaults() - require.Equal(t, "cryptosim-historical-offload", cfg.ClientID) - require.Equal(t, "none", cfg.RequiredAcks) - require.Equal(t, "snappy", cfg.Compression) - require.Equal(t, 1000, cfg.BatchSize) - require.Equal(t, 4<<20, cfg.BatchBytes) - require.NoError(t, cfg.Validate()) -} - -func TestKafkaConfigValidateRequiresBrokerAndTopic(t *testing.T) { - cfg := KafkaConfig{} - require.Error(t, cfg.Validate()) -} - -func TestKafkaConfigValidateRequiresRegionAndTLSForAWSMSKIAM(t *testing.T) { - cfg := KafkaConfig{ - Brokers: []string{"localhost:9098"}, - Topic: "historical-offload", - SASLMechanism: "aws-msk-iam", - } - cfg.ApplyDefaults() - - require.ErrorContains(t, cfg.Validate(), "tls must be enabled") - - cfg.TLSEnabled = true - require.ErrorContains(t, cfg.Validate(), "region is required") - - cfg.Region = "eu-central-1" - require.NoError(t, cfg.Validate()) -} - func TestNewKafkaStreamSupportsNilPublishAndClose(t *testing.T) { stream, err := NewKafkaStream(KafkaConfig{ Brokers: []string{"localhost:9092"}, @@ -59,9 +22,3 @@ func TestNewKafkaStreamSupportsNilPublishAndClose(t *testing.T) { require.True(t, ok) require.NoError(t, closer.Close()) } - -func TestAWSMSKIAMMechanismRequiresMetadata(t *testing.T) { - mech := &awsMSKIAMMechanism{region: "eu-central-1"} - _, _, err := mech.Start(context.Background()) - require.ErrorContains(t, err, "missing sasl metadata") -} From ba0fe3bae1b925321403962379a7e110cdadba2b Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Tue, 21 Jul 2026 13:56:07 -0400 Subject: [PATCH 4/4] Move Bigtable client library into sei-db/db_engine/bigtable Bundle the Bigtable pieces into a standalone DB library alongside the other db_engine backends: the gRPC data-plane client (ReadRows chunk assembly, MutateRows accounting, bounded retries), the frozen row-key and cell encoding, Config, and the per-RPC cost metrics on the seidb_bigtable meter. The now-redundant Bigtable prefix is dropped from names (BigtableClient -> Client, BigtableMutationRowKey -> MutationRowKey, ...); encoding behavior and metric names are unchanged and the encoding tests move with the code with assertions untouched. offload/historical keeps the Reader contract and a slim Bigtable-backed reader adapter (Get/Has/LastVersion plus mutation-row interpretation); the offload consumer sink now builds row mutations against the new package directly. Co-Authored-By: Claude Fable 5 --- sei-db/db_engine/bigtable/bigtable.go | 646 +++++++++++++++++ sei-db/db_engine/bigtable/bigtable_test.go | 177 +++++ .../bigtable}/metrics.go | 10 +- .../bigtable}/metrics_test.go | 26 +- .../state_db/ss/offload/consumer/bigtable.go | 64 +- .../ss/offload/consumer/bigtable_test.go | 66 +- .../ss/offload/historical/bigtable.go | 674 +----------------- .../ss/offload/historical/bigtable_test.go | 225 +----- 8 files changed, 957 insertions(+), 931 deletions(-) create mode 100644 sei-db/db_engine/bigtable/bigtable.go create mode 100644 sei-db/db_engine/bigtable/bigtable_test.go rename sei-db/{state_db/ss/offload/historical => db_engine/bigtable}/metrics.go (94%) rename sei-db/{state_db/ss/offload/historical => db_engine/bigtable}/metrics_test.go (63%) diff --git a/sei-db/db_engine/bigtable/bigtable.go b/sei-db/db_engine/bigtable/bigtable.go new file mode 100644 index 0000000000..7692fab086 --- /dev/null +++ b/sei-db/db_engine/bigtable/bigtable.go @@ -0,0 +1,646 @@ +// Package bigtable is a thin Bigtable data-plane client used as an MVCC +// history store: row-key encoding, ReadRows chunk assembly, MutateRows result +// accounting, bounded read retries, and per-RPC cost metrics. +package bigtable + +import ( + "context" + "encoding/binary" + "fmt" + "hash/fnv" + "io" + "regexp" + "strings" + "time" + + "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "golang.org/x/sync/errgroup" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/credentials/oauth" + "google.golang.org/grpc/keepalive" + "google.golang.org/grpc/status" +) + +const ( + DefaultFamily = "state" + DefaultShards = 256 + + mutationPrefix = byte('m') + versionPrefix = byte('v') + upgradePrefix = byte('u') + + ValueColumn = "value" + DeletedColumn = "deleted" + + // VersionBucketCount spreads monotonically increasing block-version markers + // across a bounded set of row prefixes while keeping LastVersion cheap. + VersionBucketCount = 64 +) + +// VersionBucket maps a version to its marker bucket. +func VersionBucket(version int64) int { + if version < 0 { + version = -version + } + return int(version % VersionBucketCount) +} + +const ( + bigtableEndpoint = "bigtable.googleapis.com:443" + + defaultReadWorkers = 16 + + // Transient stream failures (tablet moves, GFE restarts) surface as + // UNAVAILABLE; retry reads a bounded number of times with backoff. + readAttempts = 3 + readRetryBackoff = 100 * time.Millisecond + + maxUint16Int = 1<<16 - 1 + maxUint32Int = 1<<32 - 1 + maxInt64Uint64 = 1<<63 - 1 +) + +// Client speaks the Bigtable data protocol directly over gRPC. The official +// cloud.google.com/go/bigtable client requires a newer google.golang.org/grpc +// than the repo-wide v1.57 replace pin allows, so this package carries its +// own thin transport: ReadRows chunk assembly, MutateRows result accounting, +// bounded read retries, and per-RPC cost metrics. +type Client struct { + conn *grpc.ClientConn + data bigtablepb.BigtableClient + tableName string + appProfile string + // table is the short table name used as the metric attribute. + table string + metrics *bigtableMetrics +} + +type Cell struct { + Family string + Qualifier string + Value []byte +} + +type Row struct { + Key string + Cells []Cell +} + +type SetCell struct { + Family string + Qualifier string + TimestampMicros int64 + Value []byte +} + +type RowMutation struct { + RowKey string + SetCells []SetCell +} + +type Config struct { + ProjectID string + InstanceID string + Table string + Family string + AppProfile string + Shards int +} + +func (c *Config) ApplyDefaults() { + if c.Family == "" { + c.Family = DefaultFamily + } + if c.Shards == 0 { + c.Shards = DefaultShards + } +} + +// Configured reports whether all required connection parameters are set. +func (c Config) Configured() bool { + return strings.TrimSpace(c.ProjectID) != "" && + strings.TrimSpace(c.InstanceID) != "" && + strings.TrimSpace(c.Table) != "" +} + +func (c *Config) Validate() error { + if strings.TrimSpace(c.ProjectID) == "" { + return fmt.Errorf("bigtable project id is required") + } + if strings.TrimSpace(c.InstanceID) == "" { + return fmt.Errorf("bigtable instance id is required") + } + if strings.TrimSpace(c.Table) == "" { + return fmt.Errorf("bigtable table is required") + } + if strings.TrimSpace(c.Family) == "" { + return fmt.Errorf("bigtable family is required") + } + if c.Shards <= 0 || c.Shards > maxUint16Int { + return fmt.Errorf("bigtable shards must be between 1 and 65535") + } + return nil +} + +func NewClient(ctx context.Context, cfg Config) (*Client, error) { + cfg.ApplyDefaults() + if err := cfg.Validate(); err != nil { + return nil, err + } + perRPC, err := oauth.NewApplicationDefault(ctx, "https://www.googleapis.com/auth/bigtable.data", "https://www.googleapis.com/auth/cloud-platform") + if err != nil { + return nil, fmt.Errorf("bigtable auth: %w", err) + } + conn, err := grpc.DialContext(ctx, bigtableEndpoint, + grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), + grpc.WithPerRPCCredentials(perRPC), + // Detect silently dropped connections instead of hanging reads on the + // OS TCP timeout. + grpc.WithKeepaliveParams(keepalive.ClientParameters{ + Time: 30 * time.Second, + Timeout: 10 * time.Second, + }), + ) + if err != nil { + return nil, fmt.Errorf("open bigtable connection: %w", err) + } + return &Client{ + conn: conn, + data: bigtablepb.NewBigtableClient(conn), + tableName: fullTableName(cfg.ProjectID, cfg.InstanceID, cfg.Table), + appProfile: cfg.AppProfile, + table: cfg.Table, + metrics: newBigtableMetrics(), + }, nil +} + +type ReadRowsFunc func(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(Row) bool, qualifiers ...string) error + +type ApplyBulkFunc func(ctx context.Context, rows []RowMutation) ([]error, error) + +func (c *Client) Close() error { + if c.conn == nil { + return nil + } + return c.conn.Close() +} + +// ReadRows retries transient stream failures as long as no row has been +// delivered to the callback yet, which covers every limit-1 point read and +// version-bucket scan in this package. +func (c *Client) ReadRows(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(Row) bool, qualifiers ...string) error { + return readRowsWithRetry(ctx, c.readRowsOnce, startKey, endKey, limit, family, f, qualifiers...) +} + +func readRowsWithRetry(ctx context.Context, once ReadRowsFunc, startKey, endKey []byte, limit int64, family string, f func(Row) bool, qualifiers ...string) error { + backoff := readRetryBackoff + var delivered bool + for attempt := 1; ; attempt++ { + err := once(ctx, startKey, endKey, limit, family, func(row Row) bool { + delivered = true + return f(row) + }, qualifiers...) + if err == nil || delivered || attempt >= readAttempts || status.Code(err) != codes.Unavailable { + return err + } + if sleepErr := sleepWithContext(ctx, backoff); sleepErr != nil { + return err + } + backoff *= 2 + } +} + +func sleepWithContext(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *Client) readRowsOnce(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(Row) bool, qualifiers ...string) error { + start := time.Now() + var rowsRead, bytesRead int64 + defer func() { + c.metrics.recordRead(ctx, c.table, time.Since(start), rowsRead, bytesRead) + }() + req := &bigtablepb.ReadRowsRequest{ + TableName: c.tableName, + AppProfileId: c.appProfile, + Rows: &bigtablepb.RowSet{RowRanges: []*bigtablepb.RowRange{{ + StartKey: &bigtablepb.RowRange_StartKeyClosed{StartKeyClosed: startKey}, + EndKey: &bigtablepb.RowRange_EndKeyOpen{EndKeyOpen: endKey}, + }}}, + RowsLimit: limit, + } + if len(endKey) == 0 { + req.Rows.RowRanges[0].EndKey = nil + } + req.Filter = readFilter(family, qualifiers...) + // Cancel on return so an early callback exit tears down the server stream + // instead of leaving it to the connection finalizer. + streamCtx, cancel := context.WithCancel(ctx) + defer cancel() + stream, err := c.data.ReadRows(streamCtx, req) + if err != nil { + return err + } + var builder rowBuilder + for { + resp, err := stream.Recv() + if err == io.EOF { + return nil + } + if err != nil { + return err + } + for _, chunk := range resp.Chunks { + row, committed, err := builder.add(chunk) + if err != nil { + return err + } + if committed { + rowsRead++ + bytesRead += rowSize(row) + if !f(row) { + return nil + } + } + } + } +} + +func (c *Client) ApplyBulk(ctx context.Context, rows []RowMutation) ([]error, error) { + if len(rows) == 0 { + return nil, nil + } + start := time.Now() + var bytesWritten int64 + defer func() { + c.metrics.recordWrite(ctx, c.table, time.Since(start), int64(len(rows)), bytesWritten) + }() + entries := make([]*bigtablepb.MutateRowsRequest_Entry, 0, len(rows)) + for _, row := range rows { + bytesWritten += mutationSize(row) + entry := &bigtablepb.MutateRowsRequest_Entry{ + RowKey: []byte(row.RowKey), + Mutations: make([]*bigtablepb.Mutation, 0, len(row.SetCells)), + } + for _, cell := range row.SetCells { + entry.Mutations = append(entry.Mutations, &bigtablepb.Mutation{ + Mutation: &bigtablepb.Mutation_SetCell_{SetCell: &bigtablepb.Mutation_SetCell{ + FamilyName: cell.Family, + ColumnQualifier: []byte(cell.Qualifier), + TimestampMicros: cell.TimestampMicros, + Value: cell.Value, + }}, + }) + } + entries = append(entries, entry) + } + stream, err := c.data.MutateRows(ctx, &bigtablepb.MutateRowsRequest{ + TableName: c.tableName, + AppProfileId: c.appProfile, + Entries: entries, + }) + if err != nil { + return nil, err + } + rowErrs := make([]error, len(rows)) + seen := make([]bool, len(rows)) + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return rowErrs, err + } + for _, entry := range resp.Entries { + if entry.Index < 0 || int(entry.Index) >= len(rowErrs) { + return rowErrs, fmt.Errorf("bigtable returned invalid mutation index %d", entry.Index) + } + idx := int(entry.Index) + seen[idx] = true + if st := entry.Status; st != nil && st.Code != 0 { + rowErrs[idx] = fmt.Errorf("bigtable status %d: %s", st.Code, st.Message) + } + } + } + for i := range seen { + if !seen[i] { + rowErrs[i] = fmt.Errorf("bigtable missing mutation result") + } + } + return rowErrs, nil +} + +// LastVersion scans every version-marker bucket in parallel and returns the +// highest ingested version. +func LastVersion(ctx context.Context, readRows ReadRowsFunc) (int64, error) { + versions := make([]int64, VersionBucketCount) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(defaultReadWorkers) + for bucket := 0; bucket < VersionBucketCount; bucket++ { + bucket := bucket + g.Go(func() error { + prefix := versionRowPrefix(bucket) + var bucketVersion int64 + err := readRows(gctx, prefix, prefixEnd(prefix), 1, "", func(row Row) bool { + if version, ok := VersionFromRowKey(row.Key); ok { + bucketVersion = version + } + return false + }) + if err != nil { + return fmt.Errorf("read latest bigtable version bucket %d: %w", bucket, err) + } + versions[bucket] = bucketVersion + return nil + }) + } + if err := g.Wait(); err != nil { + return 0, err + } + var maxVersion int64 + for _, version := range versions { + if version > maxVersion { + maxVersion = version + } + } + return maxVersion, nil +} + +func mutationRowPrefixBytes(storeName string, key []byte, shards int) []byte { + shards = normalizeShards(shards) + shard := shardOf(storeName, key, shards) + prefixLen := 1 + 2 + 2 + len(storeName) + 4 + len(key) + prefix := make([]byte, prefixLen, prefixLen+8) + prefix[0] = mutationPrefix + binary.BigEndian.PutUint16(prefix[1:], shard) + binary.BigEndian.PutUint16(prefix[3:], uint16FromBoundedInt(len(storeName))) + copy(prefix[5:], storeName) + keyOffset := 5 + len(storeName) + binary.BigEndian.PutUint32(prefix[keyOffset:], uint32FromBoundedInt(len(key))) + copy(prefix[keyOffset+4:], key) + return prefix +} + +func MutationRowKey(storeName string, key []byte, version int64, shards int) string { + return string(mutationRowKeyBytes(storeName, key, version, shards)) +} + +func mutationRowKeyBytes(storeName string, key []byte, version int64, shards int) []byte { + prefix := mutationRowPrefixBytes(storeName, key, shards) + return append(prefix, invertedVersion(version)...) +} + +// MutationRowRange returns the start key of a point lookup at +// targetVersion and the exclusive end of the key's version range, hashing and +// encoding the shared row prefix only once. +func MutationRowRange(storeName string, key []byte, targetVersion int64, shards int) (start, end []byte) { + prefix := mutationRowPrefixBytes(storeName, key, shards) + end = prefixEnd(prefix) + start = append(prefix, invertedVersion(targetVersion)...) + return start, end +} + +func VersionRowKey(version int64) string { + prefix := versionRowPrefix(VersionBucket(version)) + return string(append(prefix, invertedVersion(version)...)) +} + +func UpgradeRowKey(version int64, name string) string { + key := make([]byte, 1+8+2+len(name)) + key[0] = upgradePrefix + copy(key[1:], invertedVersion(version)) + binary.BigEndian.PutUint16(key[9:], uint16FromBoundedInt(len(name))) + copy(key[11:], name) + return string(key) +} + +// VersionFromRowKey extracts the version encoded in a version-marker or +// mutation row key. +func VersionFromRowKey(rowKey string) (int64, bool) { + key := []byte(rowKey) + switch { + case len(key) >= 1+2+8 && key[0] == versionPrefix: + return decodeInvertedVersion(key[3:11]) + case len(key) >= 8 && key[0] == mutationPrefix: + return decodeInvertedVersion(key[len(key)-8:]) + default: + return 0, false + } +} + +func Timestamp(version int64) int64 { + return version * int64(time.Millisecond/time.Microsecond) +} + +type rowBuilder struct { + key []byte + family string + qualifier string + value []byte + inCell bool + cells []Cell +} + +func (b *rowBuilder) add(chunk *bigtablepb.ReadRowsResponse_CellChunk) (Row, bool, error) { + if chunk.GetResetRow() { + b.reset() + return Row{}, false, nil + } + if len(chunk.RowKey) != 0 { + b.key = append(b.key[:0], chunk.RowKey...) + b.family = "" + b.qualifier = "" + b.value = nil + b.inCell = false + b.cells = b.cells[:0] + } + if chunk.FamilyName != nil { + b.family = chunk.FamilyName.Value + } + if chunk.Qualifier != nil { + b.qualifier = string(chunk.Qualifier.Value) + b.value = b.value[:0] + b.inCell = true + } + // A non-final chunk of a split cell advertises the total value size; grow + // the buffer once instead of re-growing on every continuation chunk. + if size := int(chunk.ValueSize); size > cap(b.value) { + grown := make([]byte, len(b.value), size) + copy(grown, b.value) + b.value = grown + } + b.value = append(b.value, chunk.Value...) + if b.inCell && chunk.ValueSize == 0 { + b.cells = append(b.cells, Cell{ + Family: b.family, + Qualifier: b.qualifier, + Value: append([]byte(nil), b.value...), + }) + b.value = nil + b.inCell = false + } + if !chunk.GetCommitRow() { + return Row{}, false, nil + } + if len(b.key) == 0 { + return Row{}, false, fmt.Errorf("bigtable committed row without key") + } + row := Row{ + Key: string(append([]byte(nil), b.key...)), + Cells: append([]Cell(nil), b.cells...), + } + b.reset() + return row, true, nil +} + +func (b *rowBuilder) reset() { + b.key = nil + b.family = "" + b.qualifier = "" + b.value = nil + b.inCell = false + b.cells = nil +} + +func fullTableName(projectID, instanceID, table string) string { + return fmt.Sprintf("projects/%s/instances/%s/tables/%s", projectID, instanceID, table) +} + +func readFilter(family string, qualifiers ...string) *bigtablepb.RowFilter { + filters := make([]*bigtablepb.RowFilter, 0, 2) + if family != "" { + filters = append(filters, &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_FamilyNameRegexFilter{FamilyNameRegexFilter: regexp.QuoteMeta(family)}, + }) + } + if qualifierFilter := qualifiersFilter(qualifiers...); qualifierFilter != nil { + filters = append(filters, qualifierFilter) + } + switch len(filters) { + case 0: + return nil + case 1: + return filters[0] + default: + return &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_Chain_{Chain: &bigtablepb.RowFilter_Chain{Filters: filters}}, + } + } +} + +func qualifiersFilter(qualifiers ...string) *bigtablepb.RowFilter { + filters := make([]*bigtablepb.RowFilter, 0, len(qualifiers)) + for _, qualifier := range qualifiers { + if qualifier == "" { + continue + } + filters = append(filters, &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_ColumnQualifierRegexFilter{ + ColumnQualifierRegexFilter: []byte(regexp.QuoteMeta(qualifier)), + }, + }) + } + switch len(filters) { + case 0: + return nil + case 1: + return filters[0] + default: + return &bigtablepb.RowFilter{ + Filter: &bigtablepb.RowFilter_Interleave_{Interleave: &bigtablepb.RowFilter_Interleave{Filters: filters}}, + } + } +} + +func versionRowPrefix(bucket int) []byte { + prefix := make([]byte, 1+2) + prefix[0] = versionPrefix + binary.BigEndian.PutUint16(prefix[1:], uint16FromBoundedInt(bucket)) + return prefix +} + +func prefixEnd(prefix []byte) []byte { + end := append([]byte(nil), prefix...) + for i := len(end) - 1; i >= 0; i-- { + if end[i] != 0xff { + end[i]++ + return end[:i+1] + } + } + return nil +} + +func invertedVersion(version int64) []byte { + out := make([]byte, 8) + binary.BigEndian.PutUint64(out, ^uint64FromNonNegativeInt64(version)) + return out +} + +func decodeInvertedVersion(encoded []byte) (int64, bool) { + version := ^binary.BigEndian.Uint64(encoded) + if version > maxInt64Uint64 { + return 0, false + } + // #nosec G115 -- version is checked above to fit in int64. + return int64(version), true +} + +func shardOf(storeName string, key []byte, shards int) uint16 { + h := fnv.New32a() + _, _ = h.Write([]byte(storeName)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write(key) + return uint16FromBoundedUint32(h.Sum32() % uint32FromBoundedInt(shards)) +} + +func normalizeShards(shards int) int { + if shards <= 0 { + return DefaultShards + } + if shards > maxUint16Int { + return maxUint16Int + } + return shards +} + +func uint16FromBoundedInt(value int) uint16 { + if value < 0 || value > maxUint16Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) + } + // #nosec G115 -- value is checked above to fit in uint16. + return uint16(value) +} + +func uint32FromBoundedInt(value int) uint32 { + if value < 0 || value > maxUint32Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint32", value)) + } + // #nosec G115 -- value is checked above to fit in uint32. + return uint32(value) +} + +func uint16FromBoundedUint32(value uint32) uint16 { + if value > maxUint16Int { + panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) + } + // #nosec G115 -- value is checked above to fit in uint16. + return uint16(value) +} + +func uint64FromNonNegativeInt64(value int64) uint64 { + if value < 0 { + panic(fmt.Sprintf("bigtable version %d is negative", value)) + } + // #nosec G115 -- value is checked above to be non-negative. + return uint64(value) +} diff --git a/sei-db/db_engine/bigtable/bigtable_test.go b/sei-db/db_engine/bigtable/bigtable_test.go new file mode 100644 index 0000000000..511703bbba --- /dev/null +++ b/sei-db/db_engine/bigtable/bigtable_test.go @@ -0,0 +1,177 @@ +package bigtable + +import ( + "context" + "fmt" + "sort" + "sync" + "testing" + + "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/wrapperspb" +) + +func TestConfigDefaultsAndValidate(t *testing.T) { + cfg := Config{ + ProjectID: "project", + InstanceID: "instance", + Table: "state", + } + cfg.ApplyDefaults() + require.Equal(t, DefaultFamily, cfg.Family) + require.Equal(t, DefaultShards, cfg.Shards) + require.NoError(t, cfg.Validate()) + + missingProject := Config{InstanceID: "i", Table: "t", Family: "f"} + missingInstance := Config{ProjectID: "p", Table: "t", Family: "f"} + missingTable := Config{ProjectID: "p", InstanceID: "i", Family: "f"} + missingShards := Config{ProjectID: "p", InstanceID: "i", Table: "t", Family: "f"} + require.ErrorContains(t, missingProject.Validate(), "project") + require.ErrorContains(t, missingInstance.Validate(), "instance") + require.ErrorContains(t, missingTable.Validate(), "table") + require.ErrorContains(t, missingShards.Validate(), "shards") +} + +func TestConfigConfigured(t *testing.T) { + full := Config{ProjectID: "p", InstanceID: "i", Table: "t"} + require.True(t, full.Configured()) + + onlyProject := Config{ProjectID: "p"} + require.False(t, onlyProject.Configured()) + + require.False(t, Config{}.Configured()) +} + +func TestMutationRowKeyOrdersLatestVersionFirst(t *testing.T) { + key40 := MutationRowKey("bank", []byte("k1"), 40, 256) + key60 := MutationRowKey("bank", []byte("k1"), 60, 256) + key80 := MutationRowKey("bank", []byte("k1"), 80, 256) + keys := []string{key40, key80, key60} + sort.Strings(keys) + require.Equal(t, []string{key80, key60, key40}, keys) + + version, ok := VersionFromRowKey(key60) + require.True(t, ok) + require.Equal(t, int64(60), version) + require.NotEqual(t, + mutationRowPrefixBytes("bank", []byte("k"), 256), + mutationRowPrefixBytes("bank", []byte("k1"), 256), + ) +} + +// A cell split across chunks carries the total size in every chunk except the +// last, whose ValueSize is zero; only then may the cell be committed. +func TestRowBuilderAssemblesSplitCells(t *testing.T) { + var b rowBuilder + + row, committed, err := b.add(&bigtablepb.ReadRowsResponse_CellChunk{ + RowKey: []byte("rk"), + FamilyName: wrapperspb.String(DefaultFamily), + Qualifier: wrapperspb.Bytes([]byte(ValueColumn)), + Value: []byte("hel"), + ValueSize: 8, + }) + require.NoError(t, err) + require.False(t, committed) + require.Empty(t, row.Key) + + _, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{Value: []byte("lo "), ValueSize: 8}) + require.NoError(t, err) + require.False(t, committed) + + row, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{ + Value: []byte("bt"), + RowStatus: &bigtablepb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true}, + }) + require.NoError(t, err) + require.True(t, committed) + require.Equal(t, "rk", row.Key) + require.Len(t, row.Cells, 1) + require.Equal(t, ValueColumn, row.Cells[0].Qualifier) + require.Equal(t, []byte("hello bt"), row.Cells[0].Value) +} + +func TestReadRowsWithRetryRetriesUnavailable(t *testing.T) { + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(Row) bool, _ ...string) error { + attempts++ + if attempts < 3 { + return status.Error(codes.Unavailable, "tablet moved") + } + f(Row{Key: "rk"}) + return nil + } + var got []string + err := readRowsWithRetry(context.Background(), once, []byte("a"), []byte("b"), 1, "state", func(row Row) bool { + got = append(got, row.Key) + return false + }) + require.NoError(t, err) + require.Equal(t, 3, attempts) + require.Equal(t, []string{"rk"}, got) +} + +func TestReadRowsWithRetryStopsAfterDelivery(t *testing.T) { + // Once a row reached the callback, a retry would re-deliver it; the error + // must surface instead. + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(Row) bool, _ ...string) error { + attempts++ + f(Row{Key: "rk"}) + return status.Error(codes.Unavailable, "mid-stream drop") + } + err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(Row) bool { return true }) + require.Error(t, err) + require.Equal(t, 1, attempts) +} + +func TestReadRowsWithRetryDoesNotRetryNonRetryable(t *testing.T) { + attempts := 0 + once := func(_ context.Context, _, _ []byte, _ int64, _ string, _ func(Row) bool, _ ...string) error { + attempts++ + return status.Error(codes.NotFound, "table missing") + } + err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(Row) bool { return true }) + require.Error(t, err) + require.Equal(t, 1, attempts) +} + +func TestVersionBucket(t *testing.T) { + require.Equal(t, 0, VersionBucket(0)) + require.Equal(t, 1, VersionBucket(1)) + require.Equal(t, 0, VersionBucket(VersionBucketCount)) + require.Equal(t, 1, VersionBucket(-1)) +} + +func TestLastVersionScansBuckets(t *testing.T) { + versions := map[int]int64{ + 3: 42, + 9: 70, + } + seen := make(map[int]struct{}, VersionBucketCount) + var mu sync.Mutex + + got, err := LastVersion(context.Background(), func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(Row) bool, qualifiers ...string) error { + if len(startKey) != 3 || startKey[0] != versionPrefix { + return fmt.Errorf("unexpected start key %q", startKey) + } + if len(endKey) == 0 || limit != 1 || family != "" || len(qualifiers) != 0 { + return fmt.Errorf("unexpected scan params") + } + bucket := int(startKey[1])<<8 | int(startKey[2]) + mu.Lock() + seen[bucket] = struct{}{} + version := versions[bucket] + mu.Unlock() + if version > 0 { + f(Row{Key: VersionRowKey(version)}) + } + return nil + }) + require.NoError(t, err) + require.Equal(t, int64(70), got) + require.Len(t, seen, VersionBucketCount) +} diff --git a/sei-db/state_db/ss/offload/historical/metrics.go b/sei-db/db_engine/bigtable/metrics.go similarity index 94% rename from sei-db/state_db/ss/offload/historical/metrics.go rename to sei-db/db_engine/bigtable/metrics.go index a542121fe1..46a8987da4 100644 --- a/sei-db/state_db/ss/offload/historical/metrics.go +++ b/sei-db/db_engine/bigtable/metrics.go @@ -1,4 +1,4 @@ -package historical +package bigtable import ( "context" @@ -114,9 +114,9 @@ func (m *bigtableMetrics) recordWrite(ctx context.Context, table string, latency } } -// bigtableRowSize estimates the transferred size of a read row: the row key +// rowSize estimates the transferred size of a read row: the row key // plus each returned cell's qualifier and value. -func bigtableRowSize(row BigtableRow) int64 { +func rowSize(row Row) int64 { size := int64(len(row.Key)) for _, cell := range row.Cells { size += int64(len(cell.Qualifier)) + int64(len(cell.Value)) @@ -124,9 +124,9 @@ func bigtableRowSize(row BigtableRow) int64 { return size } -// bigtableMutationSize estimates the written size of a mutation row: the row +// mutationSize estimates the written size of a mutation row: the row // key plus each set cell's qualifier and value. -func bigtableMutationSize(row BigtableRowMutation) int64 { +func mutationSize(row RowMutation) int64 { size := int64(len(row.RowKey)) for _, cell := range row.SetCells { size += int64(len(cell.Qualifier)) + int64(len(cell.Value)) diff --git a/sei-db/state_db/ss/offload/historical/metrics_test.go b/sei-db/db_engine/bigtable/metrics_test.go similarity index 63% rename from sei-db/state_db/ss/offload/historical/metrics_test.go rename to sei-db/db_engine/bigtable/metrics_test.go index d35cfbf608..8fb9dc0ae1 100644 --- a/sei-db/state_db/ss/offload/historical/metrics_test.go +++ b/sei-db/db_engine/bigtable/metrics_test.go @@ -1,4 +1,4 @@ -package historical +package bigtable import ( "context" @@ -6,31 +6,31 @@ import ( "time" ) -func TestBigtableRowSize(t *testing.T) { - row := BigtableRow{ +func TestRowSize(t *testing.T) { + row := Row{ Key: "abc", // 3 - Cells: []BigtableCell{ + Cells: []Cell{ {Qualifier: "value", Value: []byte("hello")}, // 5 + 5 {Qualifier: "deleted", Value: []byte{0x1}}, // 7 + 1 }, } - if got, want := bigtableRowSize(row), int64(3+5+5+7+1); got != want { - t.Fatalf("bigtableRowSize() = %d, want %d", got, want) + if got, want := rowSize(row), int64(3+5+5+7+1); got != want { + t.Fatalf("rowSize() = %d, want %d", got, want) } - if got := bigtableRowSize(BigtableRow{}); got != 0 { - t.Fatalf("bigtableRowSize(empty) = %d, want 0", got) + if got := rowSize(Row{}); got != 0 { + t.Fatalf("rowSize(empty) = %d, want 0", got) } } -func TestBigtableMutationSize(t *testing.T) { - row := BigtableRowMutation{ +func TestMutationSize(t *testing.T) { + row := RowMutation{ RowKey: "row1", // 4 - SetCells: []BigtableSetCell{ + SetCells: []SetCell{ {Qualifier: "value", Value: []byte("data")}, // 5 + 4 }, } - if got, want := bigtableMutationSize(row), int64(4+5+4); got != want { - t.Fatalf("bigtableMutationSize() = %d, want %d", got, want) + if got, want := mutationSize(row), int64(4+5+4); got != want { + t.Fatalf("mutationSize() = %d, want %d", got, want) } } diff --git a/sei-db/state_db/ss/offload/consumer/bigtable.go b/sei-db/state_db/ss/offload/consumer/bigtable.go index c362d21d4e..20365e02b7 100644 --- a/sei-db/state_db/ss/offload/consumer/bigtable.go +++ b/sei-db/state_db/ss/offload/consumer/bigtable.go @@ -10,12 +10,12 @@ import ( kafkago "github.com/segmentio/kafka-go" "golang.org/x/sync/errgroup" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/bigtable" "github.com/sei-protocol/sei-chain/sei-db/proto" "github.com/sei-protocol/sei-chain/sei-db/queue/kafka" - "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" ) -type BigtableConfig = historical.BigtableConfig +type BigtableConfig = bigtable.Config const ( defaultBigtableMutationChunkRows = 1024 @@ -23,8 +23,8 @@ const ( ) type bigtableSink struct { - client *historical.BigtableClient - applyBulk historical.BigtableApplyBulkFunc + client *bigtable.Client + applyBulk bigtable.ApplyBulkFunc family string shards int bulkChunkRows int @@ -39,7 +39,7 @@ func NewBigtableSink(cfg BigtableConfig) (kafka.Sink, error) { return nil, err } ctx := context.Background() - client, err := historical.OpenBigtableClient(ctx, cfg) + client, err := bigtable.NewClient(ctx, cfg) if err != nil { return nil, err } @@ -85,7 +85,7 @@ func (s *bigtableSink) WriteBatch(ctx context.Context, msgs []kafkago.Message) e } func (s *bigtableSink) writeRecordRows(ctx context.Context, records []Record) error { - rows := make([]historical.BigtableRowMutation, 0, bigtableRowMutationCount(records)) + rows := make([]bigtable.RowMutation, 0, bigtableRowMutationCount(records)) for _, rec := range records { rows = s.appendRecordRowMutations(rows, rec.Entry) } @@ -95,7 +95,7 @@ func (s *bigtableSink) writeRecordRows(ctx context.Context, records []Record) er return s.applyRecordRowMutations(ctx, rows) } -func (s *bigtableSink) applyRecordRowMutations(ctx context.Context, rows []historical.BigtableRowMutation) error { +func (s *bigtableSink) applyRecordRowMutations(ctx context.Context, rows []bigtable.RowMutation) error { chunks := bigtableRowMutationChunks(rows, s.bulkChunkRows) g, gctx := errgroup.WithContext(ctx) g.SetLimit(s.bulkChunkWorkers) @@ -109,7 +109,7 @@ func (s *bigtableSink) applyRecordRowMutations(ctx context.Context, rows []histo return g.Wait() } -func (s *bigtableSink) appendRecordRowMutations(rows []historical.BigtableRowMutation, entry *proto.ChangelogEntry) []historical.BigtableRowMutation { +func (s *bigtableSink) appendRecordRowMutations(rows []bigtable.RowMutation, entry *proto.ChangelogEntry) []bigtable.RowMutation { for _, mutation := range compactMutations(entry) { rows = append(rows, s.mutationRow(entry.Version, mutation.storeName, mutation.pair)) } @@ -123,50 +123,50 @@ func (s *bigtableSink) appendRecordRowMutations(rows []historical.BigtableRowMut // cell for tombstones, saving a cell per delete. Readers must therefore check // the deleted column before trusting any value cell — a replayed live write // followed by a tombstone leaves both cells on the row. -func (s *bigtableSink) mutationRow(version int64, storeName string, pair *proto.KVPair) historical.BigtableRowMutation { - ts := historical.BigtableTimestamp(version) +func (s *bigtableSink) mutationRow(version int64, storeName string, pair *proto.KVPair) bigtable.RowMutation { + ts := bigtable.Timestamp(version) deleted := pair.Delete || pair.Value == nil - rowKey := historical.BigtableMutationRowKey(storeName, pair.Key, version, s.shards) + rowKey := bigtable.MutationRowKey(storeName, pair.Key, version, s.shards) if deleted { - return historical.BigtableRowMutation{ + return bigtable.RowMutation{ RowKey: rowKey, - SetCells: []historical.BigtableSetCell{{ + SetCells: []bigtable.SetCell{{ Family: s.family, - Qualifier: historical.BigtableDeletedColumn, + Qualifier: bigtable.DeletedColumn, TimestampMicros: ts, Value: boolByte(true), }}, } } - return historical.BigtableRowMutation{ + return bigtable.RowMutation{ RowKey: rowKey, - SetCells: []historical.BigtableSetCell{ - {Family: s.family, Qualifier: historical.BigtableValueColumn, TimestampMicros: ts, Value: pair.Value}, - {Family: s.family, Qualifier: historical.BigtableDeletedColumn, TimestampMicros: ts, Value: boolByte(false)}, + SetCells: []bigtable.SetCell{ + {Family: s.family, Qualifier: bigtable.ValueColumn, TimestampMicros: ts, Value: pair.Value}, + {Family: s.family, Qualifier: bigtable.DeletedColumn, TimestampMicros: ts, Value: boolByte(false)}, }, } } -func (s *bigtableSink) upgradeRow(version int64, up *proto.TreeNameUpgrade) historical.BigtableRowMutation { - ts := historical.BigtableTimestamp(version) - return historical.BigtableRowMutation{ - RowKey: historical.BigtableUpgradeRowKey(version, up.Name), - SetCells: []historical.BigtableSetCell{ +func (s *bigtableSink) upgradeRow(version int64, up *proto.TreeNameUpgrade) bigtable.RowMutation { + ts := bigtable.Timestamp(version) + return bigtable.RowMutation{ + RowKey: bigtable.UpgradeRowKey(version, up.Name), + SetCells: []bigtable.SetCell{ {Family: s.family, Qualifier: "rename_from", TimestampMicros: ts, Value: []byte(up.RenameFrom)}, - {Family: s.family, Qualifier: historical.BigtableDeletedColumn, TimestampMicros: ts, Value: boolByte(up.Delete)}, + {Family: s.family, Qualifier: bigtable.DeletedColumn, TimestampMicros: ts, Value: boolByte(up.Delete)}, }, } } func (s *bigtableSink) writeVersionMarkers(ctx context.Context, records []Record) error { - rows := make([]historical.BigtableRowMutation, 0, len(records)) + rows := make([]bigtable.RowMutation, 0, len(records)) ingestedAt := []byte(strconv.FormatInt(time.Now().UnixNano(), 10)) for _, rec := range records { version := rec.Entry.Version - ts := historical.BigtableTimestamp(version) - rows = append(rows, historical.BigtableRowMutation{ - RowKey: historical.BigtableVersionRowKey(version), - SetCells: []historical.BigtableSetCell{ + ts := bigtable.Timestamp(version) + rows = append(rows, bigtable.RowMutation{ + RowKey: bigtable.VersionRowKey(version), + SetCells: []bigtable.SetCell{ {Family: s.family, Qualifier: "topic", TimestampMicros: ts, Value: []byte(rec.Topic)}, {Family: s.family, Qualifier: "partition", TimestampMicros: ts, Value: []byte(strconv.Itoa(rec.Partition))}, {Family: s.family, Qualifier: "offset", TimestampMicros: ts, Value: []byte(strconv.FormatInt(rec.Offset, 10))}, @@ -189,7 +189,7 @@ func bigtableRowMutationCount(records []Record) int { return total } -func bigtableRowMutationChunks(rows []historical.BigtableRowMutation, maxRows int) [][]historical.BigtableRowMutation { +func bigtableRowMutationChunks(rows []bigtable.RowMutation, maxRows int) [][]bigtable.RowMutation { if len(rows) == 0 { return nil } @@ -200,7 +200,7 @@ func bigtableRowMutationChunks(rows []historical.BigtableRowMutation, maxRows in return rows[i].RowKey < rows[j].RowKey }) - chunks := make([][]historical.BigtableRowMutation, 0, (len(rows)+maxRows-1)/maxRows) + chunks := make([][]bigtable.RowMutation, 0, (len(rows)+maxRows-1)/maxRows) start := 0 startLocality := bigtableRowLocality(rows[0].RowKey) for i := 1; i < len(rows); i++ { @@ -226,7 +226,7 @@ func bigtableRowLocality(rowKey string) string { return rowKey } -func bigtableBulkError(rows []historical.BigtableRowMutation, errs []error, err error) error { +func bigtableBulkError(rows []bigtable.RowMutation, errs []error, err error) error { if err != nil { return err } diff --git a/sei-db/state_db/ss/offload/consumer/bigtable_test.go b/sei-db/state_db/ss/offload/consumer/bigtable_test.go index 28706ef822..903f13643b 100644 --- a/sei-db/state_db/ss/offload/consumer/bigtable_test.go +++ b/sei-db/state_db/ss/offload/consumer/bigtable_test.go @@ -9,8 +9,8 @@ import ( kafkago "github.com/segmentio/kafka-go" "github.com/stretchr/testify/require" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/bigtable" "github.com/sei-protocol/sei-chain/sei-db/proto" - "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/offload/historical" ) // entryMessage marshals entry into a Kafka message the sink can decode. @@ -24,10 +24,10 @@ func entryMessage(t *testing.T, entry *proto.ChangelogEntry) kafkago.Message { func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { var rows []string sink := &bigtableSink{ - family: historical.DefaultBigtableFamily, - shards: historical.DefaultBigtableShards, + family: bigtable.DefaultFamily, + shards: bigtable.DefaultShards, bulkChunkWorkers: 1, - applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + applyBulk: func(_ context.Context, mutations []bigtable.RowMutation) ([]error, error) { for _, mutation := range mutations { rows = append(rows, mutation.RowKey) } @@ -52,20 +52,20 @@ func TestBigtableSinkWritesMutationRowsAndVersionMarker(t *testing.T) { require.NoError(t, sink.WriteBatch(context.Background(), []kafkago.Message{msg})) require.Len(t, rows, 4) require.ElementsMatch(t, []string{ - historical.BigtableMutationRowKey("bank", []byte("k1"), 7, historical.DefaultBigtableShards), - historical.BigtableMutationRowKey("bank", []byte("drop"), 7, historical.DefaultBigtableShards), - historical.BigtableUpgradeRowKey(7, "new-store"), + bigtable.MutationRowKey("bank", []byte("k1"), 7, bigtable.DefaultShards), + bigtable.MutationRowKey("bank", []byte("drop"), 7, bigtable.DefaultShards), + bigtable.UpgradeRowKey(7, "new-store"), }, rows[:3]) - require.Equal(t, historical.BigtableVersionRowKey(7), rows[3]) + require.Equal(t, bigtable.VersionRowKey(7), rows[3]) } func TestBigtableSinkWriteBatchRejectsUndecodableMessage(t *testing.T) { applyBulkCalls := 0 sink := &bigtableSink{ - family: historical.DefaultBigtableFamily, - shards: historical.DefaultBigtableShards, + family: bigtable.DefaultFamily, + shards: bigtable.DefaultShards, bulkChunkWorkers: 1, - applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + applyBulk: func(_ context.Context, mutations []bigtable.RowMutation) ([]error, error) { applyBulkCalls++ return make([]error, len(mutations)), nil }, @@ -80,10 +80,10 @@ func TestBigtableSinkWriteBatchWritesRowsBeforeMarkers(t *testing.T) { var calls [][]string sink := &bigtableSink{ - family: historical.DefaultBigtableFamily, - shards: historical.DefaultBigtableShards, + family: bigtable.DefaultFamily, + shards: bigtable.DefaultShards, bulkChunkWorkers: 1, - applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + applyBulk: func(_ context.Context, mutations []bigtable.RowMutation) ([]error, error) { call := make([]string, 0, len(mutations)) for _, mutation := range mutations { call = append(call, mutation.RowKey) @@ -112,17 +112,17 @@ func TestBigtableSinkWriteBatchWritesRowsBeforeMarkers(t *testing.T) { require.NoError(t, sink.WriteBatch(context.Background(), msgs)) require.GreaterOrEqual(t, len(calls), 2) require.ElementsMatch(t, []string{ - historical.BigtableMutationRowKey("bank", []byte("k1"), 1, historical.DefaultBigtableShards), - historical.BigtableMutationRowKey("bank", []byte("k2"), 2, historical.DefaultBigtableShards), + bigtable.MutationRowKey("bank", []byte("k1"), 1, bigtable.DefaultShards), + bigtable.MutationRowKey("bank", []byte("k2"), 2, bigtable.DefaultShards), }, flattenCalls(calls[:len(calls)-1])) require.Equal(t, []string{ - historical.BigtableVersionRowKey(1), - historical.BigtableVersionRowKey(2), + bigtable.VersionRowKey(1), + bigtable.VersionRowKey(2), }, calls[len(calls)-1]) } func TestBigtableRowMutationChunksSortsAndGroupsByLocality(t *testing.T) { - rows := []historical.BigtableRowMutation{ + rows := []bigtable.RowMutation{ {RowKey: string([]byte{'u', 'z'})}, {RowKey: string([]byte{'m', 0, 2, 'b'})}, {RowKey: string([]byte{'m', 0, 1, 'c'})}, @@ -154,11 +154,11 @@ func TestBigtableRowMutationChunksSortsAndGroupsByLocality(t *testing.T) { func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { var calls [][]string sink := &bigtableSink{ - family: historical.DefaultBigtableFamily, + family: bigtable.DefaultFamily, shards: 2, bulkChunkRows: 1, bulkChunkWorkers: 1, - applyBulk: func(_ context.Context, mutations []historical.BigtableRowMutation) ([]error, error) { + applyBulk: func(_ context.Context, mutations []bigtable.RowMutation) ([]error, error) { call := make([]string, 0, len(mutations)) for _, mutation := range mutations { call = append(call, mutation.RowKey) @@ -195,18 +195,18 @@ func TestBigtableSinkWriteBatchChunksRecordRowsBeforeMarkers(t *testing.T) { require.Greater(t, len(calls), 2) markerCall := calls[len(calls)-1] require.Equal(t, []string{ - historical.BigtableVersionRowKey(1), - historical.BigtableVersionRowKey(2), + bigtable.VersionRowKey(1), + bigtable.VersionRowKey(2), }, markerCall) for _, call := range calls[:len(calls)-1] { require.Len(t, call, 1) require.NotEqual(t, markerCall, call) } require.ElementsMatch(t, []string{ - historical.BigtableMutationRowKey("bank", []byte("k1"), 1, 2), - historical.BigtableMutationRowKey("bank", []byte("k2"), 1, 2), - historical.BigtableUpgradeRowKey(1, "new-store"), - historical.BigtableMutationRowKey("bank", []byte("k3"), 2, 2), + bigtable.MutationRowKey("bank", []byte("k1"), 1, 2), + bigtable.MutationRowKey("bank", []byte("k2"), 1, 2), + bigtable.UpgradeRowKey(1, "new-store"), + bigtable.MutationRowKey("bank", []byte("k3"), 2, 2), }, flattenCalls(calls[:len(calls)-1])) } @@ -216,7 +216,7 @@ func TestBigtableSinkAppliesRecordChunksConcurrently(t *testing.T) { sink := &bigtableSink{ bulkChunkRows: 1, bulkChunkWorkers: 2, - applyBulk: func(ctx context.Context, rows []historical.BigtableRowMutation) ([]error, error) { + applyBulk: func(ctx context.Context, rows []bigtable.RowMutation) ([]error, error) { select { case started <- struct{}{}: case <-ctx.Done(): @@ -230,7 +230,7 @@ func TestBigtableSinkAppliesRecordChunksConcurrently(t *testing.T) { } }, } - rows := []historical.BigtableRowMutation{ + rows := []bigtable.RowMutation{ {RowKey: string([]byte{'m', 0, 1, 'a'})}, {RowKey: string([]byte{'m', 0, 2, 'a'})}, } @@ -253,7 +253,7 @@ func TestBigtableSinkAppliesRecordChunksConcurrently(t *testing.T) { } func TestBigtableBulkErrorValidatesMutationResultCount(t *testing.T) { - rows := []historical.BigtableRowMutation{{RowKey: "row-1"}, {RowKey: "row-2"}} + rows := []bigtable.RowMutation{{RowKey: "row-1"}, {RowKey: "row-2"}} err := bigtableBulkError(rows, []error{nil}, nil) @@ -262,7 +262,7 @@ func TestBigtableBulkErrorValidatesMutationResultCount(t *testing.T) { func TestBigtableBulkErrorWrapsRowError(t *testing.T) { rowErr := errors.New("failed") - rows := []historical.BigtableRowMutation{{RowKey: "row-1"}} + rows := []bigtable.RowMutation{{RowKey: "row-1"}} err := bigtableBulkError(rows, []error{rowErr}, nil) @@ -270,7 +270,7 @@ func TestBigtableBulkErrorWrapsRowError(t *testing.T) { require.ErrorContains(t, err, "row-1") } -func chunkRowKeys(chunks [][]historical.BigtableRowMutation) [][]string { +func chunkRowKeys(chunks [][]bigtable.RowMutation) [][]string { out := make([][]string, 0, len(chunks)) for _, chunk := range chunks { out = append(out, flattenRowMutations(chunk)) @@ -286,7 +286,7 @@ func flattenCalls(calls [][]string) []string { return out } -func flattenRowMutations(rows []historical.BigtableRowMutation) []string { +func flattenRowMutations(rows []bigtable.RowMutation) []string { out := make([]string, 0, len(rows)) for _, row := range rows { out = append(out, row.RowKey) diff --git a/sei-db/state_db/ss/offload/historical/bigtable.go b/sei-db/state_db/ss/offload/historical/bigtable.go index 64b6e4efbc..11bd3f0c63 100644 --- a/sei-db/state_db/ss/offload/historical/bigtable.go +++ b/sei-db/state_db/ss/offload/historical/bigtable.go @@ -2,151 +2,18 @@ package historical import ( "context" - "encoding/binary" "fmt" - "hash/fnv" - "io" - "regexp" - "strings" - "time" - "cloud.google.com/go/bigtable/apiv2/bigtablepb" - "golang.org/x/sync/errgroup" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials" - "google.golang.org/grpc/credentials/oauth" - "google.golang.org/grpc/keepalive" - "google.golang.org/grpc/status" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/bigtable" ) -const ( - DefaultBigtableFamily = "state" - DefaultBigtableShards = 256 - - bigtableMutationPrefix = byte('m') - bigtableVersionPrefix = byte('v') - bigtableUpgradePrefix = byte('u') - - BigtableValueColumn = "value" - BigtableDeletedColumn = "deleted" - - // VersionBucketCount spreads monotonically increasing block-version markers - // across a bounded set of row prefixes while keeping LastVersion cheap. - VersionBucketCount = 64 -) - -// VersionBucket maps a version to its marker bucket. -func VersionBucket(version int64) int { - if version < 0 { - version = -version - } - return int(version % VersionBucketCount) -} - -const ( - bigtableEndpoint = "bigtable.googleapis.com:443" - - defaultBigtableReadWorkers = 16 - - // Transient stream failures (tablet moves, GFE restarts) surface as - // UNAVAILABLE; retry reads a bounded number of times with backoff. - bigtableReadAttempts = 3 - bigtableReadRetryBackoff = 100 * time.Millisecond - - maxUint16Int = 1<<16 - 1 - maxUint32Int = 1<<32 - 1 - maxInt64Uint64 = 1<<63 - 1 -) - -// BigtableClient speaks the Bigtable data protocol directly over gRPC. The -// official cloud.google.com/go/bigtable client requires a newer -// google.golang.org/grpc than the repo-wide v1.57 replace pin allows, so this -// package carries its own thin transport: ReadRows chunk assembly, MutateRows -// result accounting, bounded read retries, and per-RPC cost metrics. -type BigtableClient struct { - conn *grpc.ClientConn - data bigtablepb.BigtableClient - tableName string - appProfile string - // table is the short table name used as the metric attribute. - table string - metrics *bigtableMetrics -} - -type BigtableCell struct { - Family string - Qualifier string - Value []byte -} - -type BigtableRow struct { - Key string - Cells []BigtableCell -} - -type BigtableSetCell struct { - Family string - Qualifier string - TimestampMicros int64 - Value []byte -} - -type BigtableRowMutation struct { - RowKey string - SetCells []BigtableSetCell -} - -type BigtableConfig struct { - ProjectID string - InstanceID string - Table string - Family string - AppProfile string - Shards int -} - -func (c *BigtableConfig) ApplyDefaults() { - if c.Family == "" { - c.Family = DefaultBigtableFamily - } - if c.Shards == 0 { - c.Shards = DefaultBigtableShards - } -} - -// Configured reports whether all required connection parameters are set. -func (c BigtableConfig) Configured() bool { - return strings.TrimSpace(c.ProjectID) != "" && - strings.TrimSpace(c.InstanceID) != "" && - strings.TrimSpace(c.Table) != "" -} - -func (c *BigtableConfig) Validate() error { - if strings.TrimSpace(c.ProjectID) == "" { - return fmt.Errorf("bigtable project id is required") - } - if strings.TrimSpace(c.InstanceID) == "" { - return fmt.Errorf("bigtable instance id is required") - } - if strings.TrimSpace(c.Table) == "" { - return fmt.Errorf("bigtable table is required") - } - if strings.TrimSpace(c.Family) == "" { - return fmt.Errorf("bigtable family is required") - } - if c.Shards <= 0 || c.Shards > maxUint16Int { - return fmt.Errorf("bigtable shards must be between 1 and 65535") - } - return nil -} - -func NewBigtableReader(cfg BigtableConfig) (Reader, error) { +// NewBigtableReader opens a Bigtable-backed historical Reader. +func NewBigtableReader(cfg bigtable.Config) (Reader, error) { cfg.ApplyDefaults() if err := cfg.Validate(); err != nil { return nil, err } - client, err := OpenBigtableClient(context.Background(), cfg) + client, err := bigtable.NewClient(context.Background(), cfg) if err != nil { return nil, err } @@ -158,205 +25,9 @@ func NewBigtableReader(cfg BigtableConfig) (Reader, error) { }, nil } -func OpenBigtableClient(ctx context.Context, cfg BigtableConfig) (*BigtableClient, error) { - cfg.ApplyDefaults() - if err := cfg.Validate(); err != nil { - return nil, err - } - perRPC, err := oauth.NewApplicationDefault(ctx, "https://www.googleapis.com/auth/bigtable.data", "https://www.googleapis.com/auth/cloud-platform") - if err != nil { - return nil, fmt.Errorf("bigtable auth: %w", err) - } - conn, err := grpc.DialContext(ctx, bigtableEndpoint, - grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), - grpc.WithPerRPCCredentials(perRPC), - // Detect silently dropped connections instead of hanging reads on the - // OS TCP timeout. - grpc.WithKeepaliveParams(keepalive.ClientParameters{ - Time: 30 * time.Second, - Timeout: 10 * time.Second, - }), - ) - if err != nil { - return nil, fmt.Errorf("open bigtable connection: %w", err) - } - return &BigtableClient{ - conn: conn, - data: bigtablepb.NewBigtableClient(conn), - tableName: bigtableTableName(cfg.ProjectID, cfg.InstanceID, cfg.Table), - appProfile: cfg.AppProfile, - table: cfg.Table, - metrics: newBigtableMetrics(), - }, nil -} - -type BigtableReadRowsFunc func(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error - -type BigtableApplyBulkFunc func(ctx context.Context, rows []BigtableRowMutation) ([]error, error) - -func (c *BigtableClient) Close() error { - if c.conn == nil { - return nil - } - return c.conn.Close() -} - -// ReadRows retries transient stream failures as long as no row has been -// delivered to the callback yet, which covers every limit-1 point read and -// version-bucket scan in this package. -func (c *BigtableClient) ReadRows(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - return readRowsWithRetry(ctx, c.readRowsOnce, startKey, endKey, limit, family, f, qualifiers...) -} - -func readRowsWithRetry(ctx context.Context, once BigtableReadRowsFunc, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - backoff := bigtableReadRetryBackoff - var delivered bool - for attempt := 1; ; attempt++ { - err := once(ctx, startKey, endKey, limit, family, func(row BigtableRow) bool { - delivered = true - return f(row) - }, qualifiers...) - if err == nil || delivered || attempt >= bigtableReadAttempts || status.Code(err) != codes.Unavailable { - return err - } - if sleepErr := sleepWithContext(ctx, backoff); sleepErr != nil { - return err - } - backoff *= 2 - } -} - -func sleepWithContext(ctx context.Context, d time.Duration) error { - timer := time.NewTimer(d) - defer timer.Stop() - select { - case <-timer.C: - return nil - case <-ctx.Done(): - return ctx.Err() - } -} - -func (c *BigtableClient) readRowsOnce(ctx context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - start := time.Now() - var rowsRead, bytesRead int64 - defer func() { - c.metrics.recordRead(ctx, c.table, time.Since(start), rowsRead, bytesRead) - }() - req := &bigtablepb.ReadRowsRequest{ - TableName: c.tableName, - AppProfileId: c.appProfile, - Rows: &bigtablepb.RowSet{RowRanges: []*bigtablepb.RowRange{{ - StartKey: &bigtablepb.RowRange_StartKeyClosed{StartKeyClosed: startKey}, - EndKey: &bigtablepb.RowRange_EndKeyOpen{EndKeyOpen: endKey}, - }}}, - RowsLimit: limit, - } - if len(endKey) == 0 { - req.Rows.RowRanges[0].EndKey = nil - } - req.Filter = bigtableReadFilter(family, qualifiers...) - // Cancel on return so an early callback exit tears down the server stream - // instead of leaving it to the connection finalizer. - streamCtx, cancel := context.WithCancel(ctx) - defer cancel() - stream, err := c.data.ReadRows(streamCtx, req) - if err != nil { - return err - } - var builder bigtableRowBuilder - for { - resp, err := stream.Recv() - if err == io.EOF { - return nil - } - if err != nil { - return err - } - for _, chunk := range resp.Chunks { - row, committed, err := builder.add(chunk) - if err != nil { - return err - } - if committed { - rowsRead++ - bytesRead += bigtableRowSize(row) - if !f(row) { - return nil - } - } - } - } -} - -func (c *BigtableClient) ApplyBulk(ctx context.Context, rows []BigtableRowMutation) ([]error, error) { - if len(rows) == 0 { - return nil, nil - } - start := time.Now() - var bytesWritten int64 - defer func() { - c.metrics.recordWrite(ctx, c.table, time.Since(start), int64(len(rows)), bytesWritten) - }() - entries := make([]*bigtablepb.MutateRowsRequest_Entry, 0, len(rows)) - for _, row := range rows { - bytesWritten += bigtableMutationSize(row) - entry := &bigtablepb.MutateRowsRequest_Entry{ - RowKey: []byte(row.RowKey), - Mutations: make([]*bigtablepb.Mutation, 0, len(row.SetCells)), - } - for _, cell := range row.SetCells { - entry.Mutations = append(entry.Mutations, &bigtablepb.Mutation{ - Mutation: &bigtablepb.Mutation_SetCell_{SetCell: &bigtablepb.Mutation_SetCell{ - FamilyName: cell.Family, - ColumnQualifier: []byte(cell.Qualifier), - TimestampMicros: cell.TimestampMicros, - Value: cell.Value, - }}, - }) - } - entries = append(entries, entry) - } - stream, err := c.data.MutateRows(ctx, &bigtablepb.MutateRowsRequest{ - TableName: c.tableName, - AppProfileId: c.appProfile, - Entries: entries, - }) - if err != nil { - return nil, err - } - rowErrs := make([]error, len(rows)) - seen := make([]bool, len(rows)) - for { - resp, err := stream.Recv() - if err == io.EOF { - break - } - if err != nil { - return rowErrs, err - } - for _, entry := range resp.Entries { - if entry.Index < 0 || int(entry.Index) >= len(rowErrs) { - return rowErrs, fmt.Errorf("bigtable returned invalid mutation index %d", entry.Index) - } - idx := int(entry.Index) - seen[idx] = true - if st := entry.Status; st != nil && st.Code != 0 { - rowErrs[idx] = fmt.Errorf("bigtable status %d: %s", st.Code, st.Message) - } - } - } - for i := range seen { - if !seen[i] { - rowErrs[i] = fmt.Errorf("bigtable missing mutation result") - } - } - return rowErrs, nil -} - type bigtableReader struct { - client *BigtableClient - readRows BigtableReadRowsFunc + client *bigtable.Client + readRows bigtable.ReadRowsFunc family string shards int } @@ -371,16 +42,16 @@ func (r *bigtableReader) Close() error { } func (r *bigtableReader) LastVersion(ctx context.Context) (int64, error) { - return BigtableLastVersion(ctx, r.readRows) + return bigtable.LastVersion(ctx, r.readRows) } func (r *bigtableReader) Has(ctx context.Context, storeName string, key []byte, targetVersion int64) (bool, error) { - start, end := bigtableMutationRowRange(storeName, key, targetVersion, r.shards) - var row BigtableRow - err := r.readRows(ctx, start, end, 1, r.family, func(r BigtableRow) bool { + start, end := bigtable.MutationRowRange(storeName, key, targetVersion, r.shards) + var row bigtable.Row + err := r.readRows(ctx, start, end, 1, r.family, func(r bigtable.Row) bool { row = r return false - }, BigtableDeletedColumn) + }, bigtable.DeletedColumn) if err != nil { return false, fmt.Errorf("bigtable has lookup: %w", err) } @@ -395,12 +66,12 @@ func (r *bigtableReader) Has(ctx context.Context, storeName string, key []byte, } func (r *bigtableReader) Get(ctx context.Context, storeName string, key []byte, targetVersion int64) (Value, error) { - start, end := bigtableMutationRowRange(storeName, key, targetVersion, r.shards) - var row BigtableRow - err := r.readRows(ctx, start, end, 1, r.family, func(r BigtableRow) bool { + start, end := bigtable.MutationRowRange(storeName, key, targetVersion, r.shards) + var row bigtable.Row + err := r.readRows(ctx, start, end, 1, r.family, func(r bigtable.Row) bool { row = r return false - }, BigtableValueColumn, BigtableDeletedColumn) + }, bigtable.ValueColumn, bigtable.DeletedColumn) if err != nil { return Value{}, fmt.Errorf("bigtable get lookup: %w", err) } @@ -410,44 +81,10 @@ func (r *bigtableReader) Get(ctx context.Context, storeName string, key []byte, return bigtableValueFromRow(row, r.family) } -func BigtableLastVersion(ctx context.Context, readRows BigtableReadRowsFunc) (int64, error) { - versions := make([]int64, VersionBucketCount) - g, gctx := errgroup.WithContext(ctx) - g.SetLimit(defaultBigtableReadWorkers) - for bucket := 0; bucket < VersionBucketCount; bucket++ { - bucket := bucket - g.Go(func() error { - prefix := bigtableVersionRowPrefix(bucket) - var bucketVersion int64 - err := readRows(gctx, prefix, bigtablePrefixEnd(prefix), 1, "", func(row BigtableRow) bool { - if version, ok := bigtableVersionFromRowKey(row.Key); ok { - bucketVersion = version - } - return false - }) - if err != nil { - return fmt.Errorf("read latest bigtable version bucket %d: %w", bucket, err) - } - versions[bucket] = bucketVersion - return nil - }) - } - if err := g.Wait(); err != nil { - return 0, err - } - var maxVersion int64 - for _, version := range versions { - if version > maxVersion { - maxVersion = version - } - } - return maxVersion, nil -} - // bigtableValueFromRow interprets a mutation row. The returned Value aliases // the row's cell buffer, which the row builder allocates per cell. -func bigtableValueFromRow(row BigtableRow, family string) (Value, error) { - version, ok := bigtableVersionFromRowKey(row.Key) +func bigtableValueFromRow(row bigtable.Row, family string) (Value, error) { + version, ok := bigtable.VersionFromRowKey(row.Key) if !ok { return Value{}, fmt.Errorf("invalid bigtable mutation row key") } @@ -458,9 +95,9 @@ func bigtableValueFromRow(row BigtableRow, family string) (Value, error) { continue } switch cell.Qualifier { - case BigtableValueColumn: + case bigtable.ValueColumn: value = cell.Value - case BigtableDeletedColumn: + case bigtable.DeletedColumn: deleted = len(cell.Value) > 0 && cell.Value[0] == 1 } } @@ -470,281 +107,14 @@ func bigtableValueFromRow(row BigtableRow, family string) (Value, error) { return Value{Bytes: value, Version: version}, nil } -func bigtableDeletedFromRow(row BigtableRow, family string) (bool, error) { - if _, ok := bigtableVersionFromRowKey(row.Key); !ok { +func bigtableDeletedFromRow(row bigtable.Row, family string) (bool, error) { + if _, ok := bigtable.VersionFromRowKey(row.Key); !ok { return false, fmt.Errorf("invalid bigtable mutation row key") } for _, cell := range row.Cells { - if cell.Family == family && cell.Qualifier == BigtableDeletedColumn { + if cell.Family == family && cell.Qualifier == bigtable.DeletedColumn { return len(cell.Value) > 0 && cell.Value[0] == 1, nil } } return false, nil } - -func bigtableMutationRowPrefixBytes(storeName string, key []byte, shards int) []byte { - shards = normalizeBigtableShards(shards) - shard := bigtableShard(storeName, key, shards) - prefixLen := 1 + 2 + 2 + len(storeName) + 4 + len(key) - prefix := make([]byte, prefixLen, prefixLen+8) - prefix[0] = bigtableMutationPrefix - binary.BigEndian.PutUint16(prefix[1:], shard) - binary.BigEndian.PutUint16(prefix[3:], uint16FromBoundedInt(len(storeName))) - copy(prefix[5:], storeName) - keyOffset := 5 + len(storeName) - binary.BigEndian.PutUint32(prefix[keyOffset:], uint32FromBoundedInt(len(key))) - copy(prefix[keyOffset+4:], key) - return prefix -} - -func BigtableMutationRowKey(storeName string, key []byte, version int64, shards int) string { - return string(bigtableMutationRowKeyBytes(storeName, key, version, shards)) -} - -func bigtableMutationRowKeyBytes(storeName string, key []byte, version int64, shards int) []byte { - prefix := bigtableMutationRowPrefixBytes(storeName, key, shards) - return append(prefix, bigtableInvertedVersion(version)...) -} - -// bigtableMutationRowRange returns the start key of a point lookup at -// targetVersion and the exclusive end of the key's version range, hashing and -// encoding the shared row prefix only once. -func bigtableMutationRowRange(storeName string, key []byte, targetVersion int64, shards int) (start, end []byte) { - prefix := bigtableMutationRowPrefixBytes(storeName, key, shards) - end = bigtablePrefixEnd(prefix) - start = append(prefix, bigtableInvertedVersion(targetVersion)...) - return start, end -} - -func BigtableVersionRowKey(version int64) string { - prefix := bigtableVersionRowPrefix(VersionBucket(version)) - return string(append(prefix, bigtableInvertedVersion(version)...)) -} - -func BigtableUpgradeRowKey(version int64, name string) string { - key := make([]byte, 1+8+2+len(name)) - key[0] = bigtableUpgradePrefix - copy(key[1:], bigtableInvertedVersion(version)) - binary.BigEndian.PutUint16(key[9:], uint16FromBoundedInt(len(name))) - copy(key[11:], name) - return string(key) -} - -func bigtableVersionFromRowKey(rowKey string) (int64, bool) { - key := []byte(rowKey) - switch { - case len(key) >= 1+2+8 && key[0] == bigtableVersionPrefix: - return bigtableDecodeInvertedVersion(key[3:11]) - case len(key) >= 8 && key[0] == bigtableMutationPrefix: - return bigtableDecodeInvertedVersion(key[len(key)-8:]) - default: - return 0, false - } -} - -func BigtableTimestamp(version int64) int64 { - return version * int64(time.Millisecond/time.Microsecond) -} - -type bigtableRowBuilder struct { - key []byte - family string - qualifier string - value []byte - inCell bool - cells []BigtableCell -} - -func (b *bigtableRowBuilder) add(chunk *bigtablepb.ReadRowsResponse_CellChunk) (BigtableRow, bool, error) { - if chunk.GetResetRow() { - b.reset() - return BigtableRow{}, false, nil - } - if len(chunk.RowKey) != 0 { - b.key = append(b.key[:0], chunk.RowKey...) - b.family = "" - b.qualifier = "" - b.value = nil - b.inCell = false - b.cells = b.cells[:0] - } - if chunk.FamilyName != nil { - b.family = chunk.FamilyName.Value - } - if chunk.Qualifier != nil { - b.qualifier = string(chunk.Qualifier.Value) - b.value = b.value[:0] - b.inCell = true - } - // A non-final chunk of a split cell advertises the total value size; grow - // the buffer once instead of re-growing on every continuation chunk. - if size := int(chunk.ValueSize); size > cap(b.value) { - grown := make([]byte, len(b.value), size) - copy(grown, b.value) - b.value = grown - } - b.value = append(b.value, chunk.Value...) - if b.inCell && chunk.ValueSize == 0 { - b.cells = append(b.cells, BigtableCell{ - Family: b.family, - Qualifier: b.qualifier, - Value: append([]byte(nil), b.value...), - }) - b.value = nil - b.inCell = false - } - if !chunk.GetCommitRow() { - return BigtableRow{}, false, nil - } - if len(b.key) == 0 { - return BigtableRow{}, false, fmt.Errorf("bigtable committed row without key") - } - row := BigtableRow{ - Key: string(append([]byte(nil), b.key...)), - Cells: append([]BigtableCell(nil), b.cells...), - } - b.reset() - return row, true, nil -} - -func (b *bigtableRowBuilder) reset() { - b.key = nil - b.family = "" - b.qualifier = "" - b.value = nil - b.inCell = false - b.cells = nil -} - -func bigtableTableName(projectID, instanceID, table string) string { - return fmt.Sprintf("projects/%s/instances/%s/tables/%s", projectID, instanceID, table) -} - -func bigtableReadFilter(family string, qualifiers ...string) *bigtablepb.RowFilter { - filters := make([]*bigtablepb.RowFilter, 0, 2) - if family != "" { - filters = append(filters, &bigtablepb.RowFilter{ - Filter: &bigtablepb.RowFilter_FamilyNameRegexFilter{FamilyNameRegexFilter: regexp.QuoteMeta(family)}, - }) - } - if qualifierFilter := bigtableQualifierFilter(qualifiers...); qualifierFilter != nil { - filters = append(filters, qualifierFilter) - } - switch len(filters) { - case 0: - return nil - case 1: - return filters[0] - default: - return &bigtablepb.RowFilter{ - Filter: &bigtablepb.RowFilter_Chain_{Chain: &bigtablepb.RowFilter_Chain{Filters: filters}}, - } - } -} - -func bigtableQualifierFilter(qualifiers ...string) *bigtablepb.RowFilter { - filters := make([]*bigtablepb.RowFilter, 0, len(qualifiers)) - for _, qualifier := range qualifiers { - if qualifier == "" { - continue - } - filters = append(filters, &bigtablepb.RowFilter{ - Filter: &bigtablepb.RowFilter_ColumnQualifierRegexFilter{ - ColumnQualifierRegexFilter: []byte(regexp.QuoteMeta(qualifier)), - }, - }) - } - switch len(filters) { - case 0: - return nil - case 1: - return filters[0] - default: - return &bigtablepb.RowFilter{ - Filter: &bigtablepb.RowFilter_Interleave_{Interleave: &bigtablepb.RowFilter_Interleave{Filters: filters}}, - } - } -} - -func bigtableVersionRowPrefix(bucket int) []byte { - prefix := make([]byte, 1+2) - prefix[0] = bigtableVersionPrefix - binary.BigEndian.PutUint16(prefix[1:], uint16FromBoundedInt(bucket)) - return prefix -} - -func bigtablePrefixEnd(prefix []byte) []byte { - end := append([]byte(nil), prefix...) - for i := len(end) - 1; i >= 0; i-- { - if end[i] != 0xff { - end[i]++ - return end[:i+1] - } - } - return nil -} - -func bigtableInvertedVersion(version int64) []byte { - out := make([]byte, 8) - binary.BigEndian.PutUint64(out, ^uint64FromNonNegativeInt64(version)) - return out -} - -func bigtableDecodeInvertedVersion(encoded []byte) (int64, bool) { - version := ^binary.BigEndian.Uint64(encoded) - if version > maxInt64Uint64 { - return 0, false - } - // #nosec G115 -- version is checked above to fit in int64. - return int64(version), true -} - -func bigtableShard(storeName string, key []byte, shards int) uint16 { - h := fnv.New32a() - _, _ = h.Write([]byte(storeName)) - _, _ = h.Write([]byte{0}) - _, _ = h.Write(key) - return uint16FromBoundedUint32(h.Sum32() % uint32FromBoundedInt(shards)) -} - -func normalizeBigtableShards(shards int) int { - if shards <= 0 { - return DefaultBigtableShards - } - if shards > maxUint16Int { - return maxUint16Int - } - return shards -} - -func uint16FromBoundedInt(value int) uint16 { - if value < 0 || value > maxUint16Int { - panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) - } - // #nosec G115 -- value is checked above to fit in uint16. - return uint16(value) -} - -func uint32FromBoundedInt(value int) uint32 { - if value < 0 || value > maxUint32Int { - panic(fmt.Sprintf("bigtable value %d exceeds uint32", value)) - } - // #nosec G115 -- value is checked above to fit in uint32. - return uint32(value) -} - -func uint16FromBoundedUint32(value uint32) uint16 { - if value > maxUint16Int { - panic(fmt.Sprintf("bigtable value %d exceeds uint16", value)) - } - // #nosec G115 -- value is checked above to fit in uint16. - return uint16(value) -} - -func uint64FromNonNegativeInt64(value int64) uint64 { - if value < 0 { - panic(fmt.Sprintf("bigtable version %d is negative", value)) - } - // #nosec G115 -- value is checked above to be non-negative. - return uint64(value) -} diff --git a/sei-db/state_db/ss/offload/historical/bigtable_test.go b/sei-db/state_db/ss/offload/historical/bigtable_test.go index fd3b2cee0d..c49e782caf 100644 --- a/sei-db/state_db/ss/offload/historical/bigtable_test.go +++ b/sei-db/state_db/ss/offload/historical/bigtable_test.go @@ -2,101 +2,48 @@ package historical import ( "context" - "fmt" - "sort" - "sync" "testing" - "cloud.google.com/go/bigtable/apiv2/bigtablepb" "github.com/stretchr/testify/require" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" - "google.golang.org/protobuf/types/known/wrapperspb" -) - -func TestBigtableConfigDefaultsAndValidate(t *testing.T) { - cfg := BigtableConfig{ - ProjectID: "project", - InstanceID: "instance", - Table: "state", - } - cfg.ApplyDefaults() - require.Equal(t, DefaultBigtableFamily, cfg.Family) - require.Equal(t, DefaultBigtableShards, cfg.Shards) - require.NoError(t, cfg.Validate()) - - missingProject := BigtableConfig{InstanceID: "i", Table: "t", Family: "f"} - missingInstance := BigtableConfig{ProjectID: "p", Table: "t", Family: "f"} - missingTable := BigtableConfig{ProjectID: "p", InstanceID: "i", Family: "f"} - missingShards := BigtableConfig{ProjectID: "p", InstanceID: "i", Table: "t", Family: "f"} - require.ErrorContains(t, missingProject.Validate(), "project") - require.ErrorContains(t, missingInstance.Validate(), "instance") - require.ErrorContains(t, missingTable.Validate(), "table") - require.ErrorContains(t, missingShards.Validate(), "shards") -} - -func TestBigtableConfigConfigured(t *testing.T) { - full := BigtableConfig{ProjectID: "p", InstanceID: "i", Table: "t"} - require.True(t, full.Configured()) - - onlyProject := BigtableConfig{ProjectID: "p"} - require.False(t, onlyProject.Configured()) - - require.False(t, BigtableConfig{}.Configured()) -} - -func TestBigtableMutationRowKeyOrdersLatestVersionFirst(t *testing.T) { - key40 := BigtableMutationRowKey("bank", []byte("k1"), 40, 256) - key60 := BigtableMutationRowKey("bank", []byte("k1"), 60, 256) - key80 := BigtableMutationRowKey("bank", []byte("k1"), 80, 256) - keys := []string{key40, key80, key60} - sort.Strings(keys) - require.Equal(t, []string{key80, key60, key40}, keys) - version, ok := bigtableVersionFromRowKey(key60) - require.True(t, ok) - require.Equal(t, int64(60), version) - require.NotEqual(t, - bigtableMutationRowPrefixBytes("bank", []byte("k"), 256), - bigtableMutationRowPrefixBytes("bank", []byte("k1"), 256), - ) -} + "github.com/sei-protocol/sei-chain/sei-db/db_engine/bigtable" +) func TestBigtableValueFromRow(t *testing.T) { - rowKey := BigtableMutationRowKey("bank", []byte("k"), 7, 256) - row := BigtableRow{ + rowKey := bigtable.MutationRowKey("bank", []byte("k"), 7, 256) + row := bigtable.Row{ Key: rowKey, - Cells: []BigtableCell{ - {Family: DefaultBigtableFamily, Qualifier: BigtableValueColumn, Value: []byte("value")}, - {Family: DefaultBigtableFamily, Qualifier: BigtableDeletedColumn, Value: []byte{0}}, + Cells: []bigtable.Cell{ + {Family: bigtable.DefaultFamily, Qualifier: bigtable.ValueColumn, Value: []byte("value")}, + {Family: bigtable.DefaultFamily, Qualifier: bigtable.DeletedColumn, Value: []byte{0}}, }, } - value, err := bigtableValueFromRow(row, DefaultBigtableFamily) + value, err := bigtableValueFromRow(row, bigtable.DefaultFamily) require.NoError(t, err) require.Equal(t, []byte("value"), value.Bytes) require.Equal(t, int64(7), value.Version) row.Cells[1].Value = []byte{1} - _, err = bigtableValueFromRow(row, DefaultBigtableFamily) + _, err = bigtableValueFromRow(row, bigtable.DefaultFamily) require.ErrorIs(t, err, ErrNotFound) } func TestBigtableReaderGetUsesMVCCRange(t *testing.T) { - wantRow := BigtableMutationRowKey("bank", []byte("k"), 40, 256) + wantRow := bigtable.MutationRowKey("bank", []byte("k"), 40, 256) reader := &bigtableReader{ - family: DefaultBigtableFamily, + family: bigtable.DefaultFamily, shards: 256, - readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - require.Equal(t, []byte(BigtableMutationRowKey("bank", []byte("k"), 60, 256)), startKey) + readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(bigtable.Row) bool, qualifiers ...string) error { + require.Equal(t, []byte(bigtable.MutationRowKey("bank", []byte("k"), 60, 256)), startKey) require.NotEmpty(t, endKey) require.Equal(t, int64(1), limit) - require.Equal(t, DefaultBigtableFamily, family) - require.Equal(t, []string{BigtableValueColumn, BigtableDeletedColumn}, qualifiers) - f(BigtableRow{ + require.Equal(t, bigtable.DefaultFamily, family) + require.Equal(t, []string{bigtable.ValueColumn, bigtable.DeletedColumn}, qualifiers) + f(bigtable.Row{ Key: wantRow, - Cells: []BigtableCell{ - {Family: DefaultBigtableFamily, Qualifier: BigtableValueColumn, Value: []byte("v40")}, - {Family: DefaultBigtableFamily, Qualifier: BigtableDeletedColumn, Value: []byte{0}}, + Cells: []bigtable.Cell{ + {Family: bigtable.DefaultFamily, Qualifier: bigtable.ValueColumn, Value: []byte("v40")}, + {Family: bigtable.DefaultFamily, Qualifier: bigtable.DeletedColumn, Value: []byte{0}}, }, }) return nil @@ -110,21 +57,21 @@ func TestBigtableReaderGetUsesMVCCRange(t *testing.T) { } func TestBigtableReaderHasOnlyReadsDeletedColumn(t *testing.T) { - wantRow := BigtableMutationRowKey("bank", []byte("k"), 40, 256) + wantRow := bigtable.MutationRowKey("bank", []byte("k"), 40, 256) reader := &bigtableReader{ - family: DefaultBigtableFamily, + family: bigtable.DefaultFamily, shards: 256, - readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - require.Equal(t, []byte(BigtableMutationRowKey("bank", []byte("k"), 60, 256)), startKey) + readRows: func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(bigtable.Row) bool, qualifiers ...string) error { + require.Equal(t, []byte(bigtable.MutationRowKey("bank", []byte("k"), 60, 256)), startKey) require.NotEmpty(t, endKey) require.Equal(t, int64(1), limit) - require.Equal(t, DefaultBigtableFamily, family) - require.Equal(t, []string{BigtableDeletedColumn}, qualifiers) - f(BigtableRow{ + require.Equal(t, bigtable.DefaultFamily, family) + require.Equal(t, []string{bigtable.DeletedColumn}, qualifiers) + f(bigtable.Row{ Key: wantRow, - Cells: []BigtableCell{{ - Family: DefaultBigtableFamily, - Qualifier: BigtableDeletedColumn, + Cells: []bigtable.Cell{{ + Family: bigtable.DefaultFamily, + Qualifier: bigtable.DeletedColumn, Value: []byte{0}, }}, }) @@ -136,117 +83,3 @@ func TestBigtableReaderHasOnlyReadsDeletedColumn(t *testing.T) { require.NoError(t, err) require.True(t, ok) } - -// A cell split across chunks carries the total size in every chunk except the -// last, whose ValueSize is zero; only then may the cell be committed. -func TestBigtableRowBuilderAssemblesSplitCells(t *testing.T) { - var b bigtableRowBuilder - - row, committed, err := b.add(&bigtablepb.ReadRowsResponse_CellChunk{ - RowKey: []byte("rk"), - FamilyName: wrapperspb.String(DefaultBigtableFamily), - Qualifier: wrapperspb.Bytes([]byte(BigtableValueColumn)), - Value: []byte("hel"), - ValueSize: 8, - }) - require.NoError(t, err) - require.False(t, committed) - require.Empty(t, row.Key) - - _, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{Value: []byte("lo "), ValueSize: 8}) - require.NoError(t, err) - require.False(t, committed) - - row, committed, err = b.add(&bigtablepb.ReadRowsResponse_CellChunk{ - Value: []byte("bt"), - RowStatus: &bigtablepb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true}, - }) - require.NoError(t, err) - require.True(t, committed) - require.Equal(t, "rk", row.Key) - require.Len(t, row.Cells, 1) - require.Equal(t, BigtableValueColumn, row.Cells[0].Qualifier) - require.Equal(t, []byte("hello bt"), row.Cells[0].Value) -} - -func TestReadRowsWithRetryRetriesUnavailable(t *testing.T) { - attempts := 0 - once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(BigtableRow) bool, _ ...string) error { - attempts++ - if attempts < 3 { - return status.Error(codes.Unavailable, "tablet moved") - } - f(BigtableRow{Key: "rk"}) - return nil - } - var got []string - err := readRowsWithRetry(context.Background(), once, []byte("a"), []byte("b"), 1, "state", func(row BigtableRow) bool { - got = append(got, row.Key) - return false - }) - require.NoError(t, err) - require.Equal(t, 3, attempts) - require.Equal(t, []string{"rk"}, got) -} - -func TestReadRowsWithRetryStopsAfterDelivery(t *testing.T) { - // Once a row reached the callback, a retry would re-deliver it; the error - // must surface instead. - attempts := 0 - once := func(_ context.Context, _, _ []byte, _ int64, _ string, f func(BigtableRow) bool, _ ...string) error { - attempts++ - f(BigtableRow{Key: "rk"}) - return status.Error(codes.Unavailable, "mid-stream drop") - } - err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(BigtableRow) bool { return true }) - require.Error(t, err) - require.Equal(t, 1, attempts) -} - -func TestReadRowsWithRetryDoesNotRetryNonRetryable(t *testing.T) { - attempts := 0 - once := func(_ context.Context, _, _ []byte, _ int64, _ string, _ func(BigtableRow) bool, _ ...string) error { - attempts++ - return status.Error(codes.NotFound, "table missing") - } - err := readRowsWithRetry(context.Background(), once, nil, nil, 0, "", func(BigtableRow) bool { return true }) - require.Error(t, err) - require.Equal(t, 1, attempts) -} - -func TestVersionBucket(t *testing.T) { - require.Equal(t, 0, VersionBucket(0)) - require.Equal(t, 1, VersionBucket(1)) - require.Equal(t, 0, VersionBucket(VersionBucketCount)) - require.Equal(t, 1, VersionBucket(-1)) -} - -func TestBigtableLastVersionScansBuckets(t *testing.T) { - versions := map[int]int64{ - 3: 42, - 9: 70, - } - seen := make(map[int]struct{}, VersionBucketCount) - var mu sync.Mutex - - got, err := BigtableLastVersion(context.Background(), func(_ context.Context, startKey, endKey []byte, limit int64, family string, f func(BigtableRow) bool, qualifiers ...string) error { - if len(startKey) != 3 || startKey[0] != bigtableVersionPrefix { - return fmt.Errorf("unexpected start key %q", startKey) - } - if len(endKey) == 0 || limit != 1 || family != "" || len(qualifiers) != 0 { - return fmt.Errorf("unexpected scan params") - } - bucket := int(startKey[1])<<8 | int(startKey[2]) - mu.Lock() - seen[bucket] = struct{}{} - version := versions[bucket] - mu.Unlock() - if version > 0 { - f(BigtableRow{Key: BigtableVersionRowKey(version)}) - } - return nil - }) - require.NoError(t, err) - require.Equal(t, int64(70), got) - require.Len(t, seen, VersionBucketCount) -}