From 0c5d3dd587a614b0c681030b4cec372a21dc6463 Mon Sep 17 00:00:00 2001 From: kbhat1 Date: Fri, 17 Jul 2026 17:15:52 -0400 Subject: [PATCH] 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: