feat: boilerplate accelerator code Change-Id: I48b4a71668d4f444f5aa4961688c83876f86a64f
diff --git a/bigtable/accelerator/accelerator_core.go b/bigtable/accelerator/accelerator_core.go new file mode 100644 index 0000000..27752fa --- /dev/null +++ b/bigtable/accelerator/accelerator_core.go
@@ -0,0 +1,147 @@ +package accelerator + +import ( + "context" + + "accelerator/adapters" + "accelerator/metrics" + "accelerator/resourcemanager" + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" +) + +// SessionPool is a placeholder for the future per-(resource, method) session +// pool cached by the resource manager. Its ReadRow/MutateRow stubs return +// Unimplemented until the Jetstream transport is wired; the rest of the +// translation pipeline (dispatch, pool lookup) is built against this type +// so swapping in the real implementation is a localized change. +type SessionPool struct{} + +// Close releases resources owned by the pool. No-op for the placeholder. +func (p *SessionPool) Close() error { return nil } + +func (p *SessionPool) ReadRow(ctx context.Context, req *v2pb.SessionReadRowRequest) (*v2pb.SessionReadRowResponse, error) { + return nil, status.Error(codes.Unimplemented, "SessionPool.ReadRow not wired") +} + +func (p *SessionPool) MutateRow(ctx context.Context, req *v2pb.SessionMutateRowRequest) (*v2pb.SessionMutateRowResponse, error) { + return nil, status.Error(codes.Unimplemented, "SessionPool.MutateRow not wired") +} + +// Ensure AcceleratorChannel implements grpc.ClientConnInterface. +var _ grpc.ClientConnInterface = (*AcceleratorChannel)(nil) + +// AcceleratorChannel is the single in-process channel used by both the native +// Go client and the daemon's gRPC service via grpc.ClientConnInterface. +type AcceleratorChannel struct { + pools *resourcemanager.PoolCache[*SessionPool] + recorder *metrics.MetricsRecorder +} + +// NewAcceleratorChannel constructs an AcceleratorChannel with its pool cache +// and metrics recorder wired internally. This is the public factory consumed +// by both the native Go client and the daemon binary. +func NewAcceleratorChannel() *AcceleratorChannel { + factory := func(resource, method string) (*SessionPool, error) { + return &SessionPool{}, nil + } + return &AcceleratorChannel{ + pools: resourcemanager.NewPoolCache[*SessionPool](resourcemanager.DefaultPoolCacheSize, factory), + recorder: &metrics.MetricsRecorder{}, + } +} + +// Invoke implements grpc.ClientConnInterface by dispatching on the V2 method +// name to a per-method impl helper. +func (c *AcceleratorChannel) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error { + switch method { + case v2pb.Bigtable_MutateRow_FullMethodName: + return c.mutateRowImpl(ctx, args, reply) + default: + return status.Errorf(codes.Unimplemented, "method %s not implemented", method) + } +} + + +func (c *AcceleratorChannel) mutateRowImpl(ctx context.Context, args interface{}, reply interface{}) error { + reqV2, ok := args.(*v2pb.MutateRowRequest) + if !ok { + return status.Errorf(codes.Internal, "unexpected request type: %T", args) + } + respV2, ok := reply.(*v2pb.MutateRowResponse) + if !ok { + return status.Errorf(codes.Internal, "unexpected response type: %T", reply) + } + + reqAdapter := &adapters.MutateRowRequestAdapter{} + reqJS, err := reqAdapter.Adapt(reqV2) + if err != nil { + return err + } + resourceName, err := reqAdapter.ExtractResource(reqV2) + if err != nil { + return err + } + + pool, release, err := c.pools.GetOrOpen(resourceName, "MutateRow") + if err != nil { + return err + } + defer release() + respJS, err := pool.MutateRow(ctx, reqJS) + if err != nil { + return err + } + + respAdapter := &adapters.MutateRowResponseAdapter{} + adaptedResp, err := respAdapter.Adapt(respJS) + if err != nil { + return err + } + + proto.Reset(respV2) + proto.Merge(respV2, adaptedResp) + return nil +} + +func (c *AcceleratorChannel) readRowsImpl(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return &readRowsClientStream{ctx: ctx}, nil +} + +// NewStream implements grpc.ClientConnInterface. +func (c *AcceleratorChannel) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + switch method { + case v2pb.Bigtable_ReadRows_FullMethodName: + return c.readRowsImpl(ctx, desc, method, opts...) + default: + return nil, status.Errorf(codes.Unimplemented, "method %s not implemented", method) + } +} + +type readRowsClientStream struct { + ctx context.Context +} + +func (s *readRowsClientStream) Header() (metadata.MD, error) { return nil, nil } +func (s *readRowsClientStream) Trailer() metadata.MD { return nil } +func (s *readRowsClientStream) CloseSend() error { return nil } +func (s *readRowsClientStream) Context() context.Context { return s.ctx } +func (s *readRowsClientStream) SendMsg(m any) error { return status.Error(codes.Unimplemented, "SendMsg not implemented") } +func (s *readRowsClientStream) RecvMsg(m any) error { + return status.Error(codes.Unimplemented, "RecvMsg not implemented") +} + +// Close releases resources held by the channel, draining any pools owned by +// the resource manager. Callers (notably AcceleratorServer.Stop) should +// invoke this after the gRPC server has finished draining in-flight RPCs so +// pools are not torn down while requests are still using them. +func (c *AcceleratorChannel) Close() error { + if c.pools == nil { + return nil + } + return c.pools.Close() +}
diff --git a/bigtable/accelerator/accelerator_core_test.go b/bigtable/accelerator/accelerator_core_test.go new file mode 100644 index 0000000..9b1f5cf --- /dev/null +++ b/bigtable/accelerator/accelerator_core_test.go
@@ -0,0 +1,61 @@ +package accelerator + +import ( + "context" + "testing" + + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestNewAcceleratorChannel(t *testing.T) { + channel := NewAcceleratorChannel() + if channel == nil { + t.Fatal("NewAcceleratorChannel returned nil") + } +} + +// Locks in the contract the future V2 shim depends on: until SessionPool's +// ReadRow/MutateRow stubs are replaced, the full translation pipeline +// runs and then surfaces codes.Unimplemented from the pool stub. +func TestInvoke_ReturnsUnimplemented(t *testing.T) { + channel := NewAcceleratorChannel() + ctx := context.Background() + + mutateReq := &v2pb.MutateRowRequest{ + TableName: "projects/p/instances/i/tables/t", + RowKey: []byte("k"), + } + err := channel.Invoke(ctx, "/google.bigtable.v2.Bigtable/MutateRow", mutateReq, &v2pb.MutateRowResponse{}) + if status.Code(err) != codes.Unimplemented { + t.Errorf("Invoke(MutateRow) code = %v; want Unimplemented", status.Code(err)) + } + + err = channel.Invoke(ctx, "unknown", nil, nil) + if status.Code(err) != codes.Unimplemented { + t.Errorf("Invoke(unknown) code = %v; want Unimplemented", status.Code(err)) + } +} + +func TestNewStream(t *testing.T) { + channel := NewAcceleratorChannel() + ctx := context.Background() + + stream, err := channel.NewStream(ctx, nil, "/google.bigtable.v2.Bigtable/ReadRows") + if err != nil { + t.Fatalf("NewStream(ReadRows) failed: %v", err) + } + if stream == nil { + t.Fatal("NewStream(ReadRows) returned nil stream") + } + err = stream.RecvMsg(nil) + if status.Code(err) != codes.Unimplemented { + t.Errorf("stream.RecvMsg() code = %v; want Unimplemented", status.Code(err)) + } + + _, err = channel.NewStream(ctx, nil, "unknown") + if status.Code(err) != codes.Unimplemented { + t.Errorf("NewStream(unknown) code = %v; want Unimplemented", status.Code(err)) + } +}
diff --git a/bigtable/accelerator/adapters/common.go b/bigtable/accelerator/adapters/common.go new file mode 100644 index 0000000..4383ab3 --- /dev/null +++ b/bigtable/accelerator/adapters/common.go
@@ -0,0 +1,20 @@ +package adapters + +// Adapter defines a generic interface for adapting one type to another. +type Adapter[From any, To any] interface { + Adapt(from From) (To, error) +} + +// RequestAdapter represents a specialized adapter for request routing. +type RequestAdapter[From any, To any] interface { + Adapter[From, To] + ExtractResource(from From) (string, error) +} + +// Default request and response adapter singletons. +var ( + DefaultReadRowRequestAdapter = &ReadRowRequestAdapter{} + DefaultReadRowResponseAdapter = &ReadRowResponseAdapter{} + DefaultMutateRowRequestAdapter = &MutateRowRequestAdapter{} + DefaultMutateRowResponseAdapter = &MutateRowResponseAdapter{} +)
diff --git a/bigtable/accelerator/adapters/common_test.go b/bigtable/accelerator/adapters/common_test.go new file mode 100644 index 0000000..5ae6521 --- /dev/null +++ b/bigtable/accelerator/adapters/common_test.go
@@ -0,0 +1,18 @@ +package adapters + +import "testing" + +func TestAdaptersExist(t *testing.T) { + if DefaultReadRowRequestAdapter == nil { + t.Error("expected DefaultReadRowRequestAdapter to be non-nil") + } + if DefaultReadRowResponseAdapter == nil { + t.Error("expected DefaultReadRowResponseAdapter to be non-nil") + } + if DefaultMutateRowRequestAdapter == nil { + t.Error("expected DefaultMutateRowRequestAdapter to be non-nil") + } + if DefaultMutateRowResponseAdapter == nil { + t.Error("expected DefaultMutateRowResponseAdapter to be non-nil") + } +}
diff --git a/bigtable/accelerator/adapters/mutate_row.go b/bigtable/accelerator/adapters/mutate_row.go new file mode 100644 index 0000000..b38e3cc --- /dev/null +++ b/bigtable/accelerator/adapters/mutate_row.go
@@ -0,0 +1,38 @@ +package adapters + +import ( + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// MutateRowRequestAdapter adapts V2 MutateRowRequest to SessionMutateRowRequest. +type MutateRowRequestAdapter struct{} + +func (a *MutateRowRequestAdapter) Adapt(from *v2pb.MutateRowRequest) (*v2pb.SessionMutateRowRequest, error) { + if from == nil { + return nil, nil + } + return &v2pb.SessionMutateRowRequest{ + Key: from.RowKey, + Mutations: from.Mutations, + }, nil +} + +func (a *MutateRowRequestAdapter) ExtractResource(from *v2pb.MutateRowRequest) (string, error) { + if from == nil { + return "", status.Errorf(codes.InvalidArgument, "request is nil") + } + return from.TableName, nil +} + +// MutateRowResponseAdapter adapts SessionMutateRowResponse to MutateRowResponse. +type MutateRowResponseAdapter struct{} + +func (a *MutateRowResponseAdapter) Adapt(from *v2pb.SessionMutateRowResponse) (*v2pb.MutateRowResponse, error) { + if from == nil { + return nil, nil + } + // Bare minimum scaffold. + return &v2pb.MutateRowResponse{}, nil +}
diff --git a/bigtable/accelerator/adapters/mutate_row_test.go b/bigtable/accelerator/adapters/mutate_row_test.go new file mode 100644 index 0000000..d7883d6 --- /dev/null +++ b/bigtable/accelerator/adapters/mutate_row_test.go
@@ -0,0 +1,69 @@ +package adapters + +import ( + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "testing" +) + +func TestMutateRowRequestAdapter(t *testing.T) { + reqAdapter := &MutateRowRequestAdapter{} + v2Req := &v2pb.MutateRowRequest{ + TableName: "projects/p1/instances/i1/tables/t1", + RowKey: []byte("row-key"), + Mutations: []*v2pb.Mutation{ + { + Mutation: &v2pb.Mutation_SetCell_{ + SetCell: &v2pb.Mutation_SetCell{ + FamilyName: "fam", + ColumnQualifier: []byte("qual"), + Value: []byte("val"), + TimestampMicros: 1000, + }, + }, + }, + }, + } + + jsReq, err := reqAdapter.Adapt(v2Req) + if err != nil { + t.Fatalf("Adapt failed: %v", err) + } + + if string(jsReq.Key) != "row-key" { + t.Errorf("expected key 'row-key', got %s", string(jsReq.Key)) + } + + if len(jsReq.Mutations) != 1 { + t.Fatalf("expected 1 mutation, got %d", len(jsReq.Mutations)) + } + + setCell := jsReq.Mutations[0].GetSetCell() + if setCell == nil { + t.Fatal("expected SetCell mutation") + } + + if setCell.FamilyName != "fam" || string(setCell.ColumnQualifier) != "qual" || string(setCell.Value) != "val" { + t.Errorf("unexpected set cell content: %+v", setCell) + } + + res, err := reqAdapter.ExtractResource(v2Req) + if err != nil { + t.Fatalf("ExtractResource failed: %v", err) + } + if res != "projects/p1/instances/i1/tables/t1" { + t.Errorf("expected resource projects/p1/instances/i1/tables/t1, got %s", res) + } +} + +func TestMutateRowResponseAdapter(t *testing.T) { + resAdapter := &MutateRowResponseAdapter{} + jsRes := &v2pb.SessionMutateRowResponse{} + v2Res, err := resAdapter.Adapt(jsRes) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if v2Res == nil { + t.Fatal("expected non-nil v2Res") + } +}
diff --git a/bigtable/accelerator/adapters/read_row.go b/bigtable/accelerator/adapters/read_row.go new file mode 100644 index 0000000..89f9661 --- /dev/null +++ b/bigtable/accelerator/adapters/read_row.go
@@ -0,0 +1,40 @@ +package adapters + +import ( + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// ReadRowRequestAdapter adapts V2 ReadRowsRequest to SessionReadRowRequest. +type ReadRowRequestAdapter struct{} + +func (a *ReadRowRequestAdapter) Adapt(from *v2pb.ReadRowsRequest) (*v2pb.SessionReadRowRequest, error) { + if from == nil { + return nil, nil + } + req := &v2pb.SessionReadRowRequest{} + if from.Rows != nil && len(from.Rows.RowKeys) > 0 { + req.Key = from.Rows.RowKeys[0] + } + req.Filter = from.Filter + return req, nil +} + +func (a *ReadRowRequestAdapter) ExtractResource(from *v2pb.ReadRowsRequest) (string, error) { + if from == nil { + return "", status.Errorf(codes.InvalidArgument, "request is nil") + } + return from.TableName, nil +} + +// ReadRowResponseAdapter adapts SessionReadRowResponse to ReadRowsResponse. +type ReadRowResponseAdapter struct{} + +func (a *ReadRowResponseAdapter) Adapt(from *v2pb.SessionReadRowResponse) (*v2pb.ReadRowsResponse, error) { + if from == nil { + return nil, nil + } + // Bare minimum scaffold. + return &v2pb.ReadRowsResponse{}, nil +}
diff --git a/bigtable/accelerator/adapters/read_row_test.go b/bigtable/accelerator/adapters/read_row_test.go new file mode 100644 index 0000000..089c392 --- /dev/null +++ b/bigtable/accelerator/adapters/read_row_test.go
@@ -0,0 +1,56 @@ +package adapters + +import ( + v2pb "cloud.google.com/go/bigtable/apiv2/bigtablepb" + "testing" +) + +func TestReadRowRequestAdapter(t *testing.T) { + reqAdapter := &ReadRowRequestAdapter{} + v2Req := &v2pb.ReadRowsRequest{ + TableName: "projects/p1/instances/i1/tables/t1", + Rows: &v2pb.RowSet{ + RowKeys: [][]byte{[]byte("test-key")}, + }, + Filter: &v2pb.RowFilter{ + Filter: &v2pb.RowFilter_FamilyNameRegexFilter{FamilyNameRegexFilter: "family-regex"}, + }, + } + jsReq, err := reqAdapter.Adapt(v2Req) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if jsReq == nil { + t.Fatal("expected non-nil jsReq") + } + + if string(jsReq.Key) != "test-key" { + t.Errorf("expected Key test-key, got %s", string(jsReq.Key)) + } + + if jsReq.Filter.GetFamilyNameRegexFilter() != "family-regex" { + t.Errorf("expected Filter family-regex, got %s", jsReq.Filter.GetFamilyNameRegexFilter()) + } + + res, err := reqAdapter.ExtractResource(v2Req) + if err != nil { + t.Fatalf("ExtractResource failed: %v", err) + } + if res != "projects/p1/instances/i1/tables/t1" { + t.Errorf("expected resource projects/p1/instances/i1/tables/t1, got %s", res) + } +} + +func TestReadRowResponseAdapter(t *testing.T) { + resAdapter := &ReadRowResponseAdapter{} + jsRes := &v2pb.SessionReadRowResponse{} + v2Res, err := resAdapter.Adapt(jsRes) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if v2Res == nil { + t.Fatal("expected non-nil v2Res") + } +}
diff --git a/bigtable/accelerator/go.mod b/bigtable/accelerator/go.mod new file mode 100644 index 0000000..7baca9c --- /dev/null +++ b/bigtable/accelerator/go.mod
@@ -0,0 +1,21 @@ +module accelerator + +go 1.25.8 + +require ( + cloud.google.com/go/bigtable v1.47.0 + github.com/hashicorp/golang-lru/v2 v2.0.7 + google.golang.org/grpc v1.81.1 + google.golang.org/protobuf v1.36.11 +) + +require ( + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect +) + +replace cloud.google.com/go/bigtable => ../
diff --git a/bigtable/accelerator/go.sum b/bigtable/accelerator/go.sum new file mode 100644 index 0000000..3ab303b --- /dev/null +++ b/bigtable/accelerator/go.sum
@@ -0,0 +1,44 @@ +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= +go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= +go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= +go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= +go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= +go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94 h1:YJjbgu+dkp5kUJLfpMyCLfBIWZb/FcJyuLeo1gVBOuo= +google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94/go.mod h1:RRHjglSYABVCWpQ7USCpdfhcd9t4PkajvVwyynZizTc= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= +google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
diff --git a/bigtable/accelerator/metrics/metrics.go b/bigtable/accelerator/metrics/metrics.go new file mode 100644 index 0000000..eaa9671 --- /dev/null +++ b/bigtable/accelerator/metrics/metrics.go
@@ -0,0 +1,4 @@ +package metrics + +// MetricsRecorder is a dummy struct for recording metrics. +type MetricsRecorder struct{}
diff --git a/bigtable/accelerator/resourcemanager/resourcemanager.go b/bigtable/accelerator/resourcemanager/resourcemanager.go new file mode 100644 index 0000000..aafd518 --- /dev/null +++ b/bigtable/accelerator/resourcemanager/resourcemanager.go
@@ -0,0 +1,164 @@ +package resourcemanager + +import ( + "io" + "log" + "sync" + "sync/atomic" + "time" + + "github.com/hashicorp/golang-lru/v2/expirable" +) + +// DefaultPoolCacheSize is the default capacity used when a non-positive +// capacity is supplied to NewPoolCache. +const DefaultPoolCacheSize = 16 + +// DefaultPoolCacheTTL is the per-entry idle expiry. Entries that have not +// been fetched via GetOrOpen within this window are evicted; each GetOrOpen +// refreshes the entry's TTL (sliding expiry). +const DefaultPoolCacheTTL = 5 * time.Minute + +// PoolFactory constructs a pool for a (resource, method) pair on cache miss. +// The accelerator calls this in PoolCache.GetOrOpen. +type PoolFactory[V io.Closer] func(resource, method string) (V, error) + +// PoolKey is the composite cache key: one entry per (resource, method) pair. +// Exported so callers can reason about cache identity if needed. +type PoolKey struct { + Resource string + Method string +} + +// entry wraps a cached pool with the bookkeeping needed for safe concurrent +// borrow/release across LRU and TTL eviction. Refs counts outstanding handles +// returned by GetOrOpen; evicted flips true once the cache has detached the +// entry (LRU pressure, TTL expiry, or Close). The underlying pool's Close +// runs exactly once, fired by whichever observer (release or eviction) last +// sees refs==0 && evicted==true. +type entry[V io.Closer] struct { + key PoolKey + pool V + refs atomic.Int32 + evicted atomic.Bool + closer sync.Once +} + +func (e *entry[V]) close() error { + var err error + e.closer.Do(func() { + err = e.pool.Close() + }) + return err +} + +func (e *entry[V]) maybeClose() { + if e.evicted.Load() && e.refs.Load() == 0 { + if err := e.close(); err != nil { + log.Printf("resourcemanager: Close on evicted pool %+v: %v", e.key, err) + } + } +} + +// PoolCache caches one pool per (resource, method) pair, bounded by an LRU +// and a per-entry idle TTL. Evicted pools are closed once their last +// outstanding borrowed handle is released. +type PoolCache[V io.Closer] struct { + factory PoolFactory[V] + mu sync.Mutex + cache *expirable.LRU[PoolKey, *entry[V]] +} + +// NewPoolCache returns a PoolCache that constructs pools on miss via factory, +// caches up to capacity entries, and expires idle entries after +// DefaultPoolCacheTTL. A capacity <= 0 falls back to DefaultPoolCacheSize. +func NewPoolCache[V io.Closer](capacity int, factory PoolFactory[V]) *PoolCache[V] { + return newPoolCacheWithTTL(capacity, DefaultPoolCacheTTL, factory) +} + +// newPoolCacheWithTTL is the constructor backing NewPoolCache; tests use it +// directly with a short TTL. +func newPoolCacheWithTTL[V io.Closer](capacity int, ttl time.Duration, factory PoolFactory[V]) *PoolCache[V] { + if capacity <= 0 { + capacity = DefaultPoolCacheSize + } + m := &PoolCache[V]{factory: factory} + m.cache = expirable.NewLRU[PoolKey, *entry[V]](capacity, onCacheEvict[V], ttl) + return m +} + +// onCacheEvict fires for both LRU and TTL eviction. It runs under the cache's +// internal lock (and, for LRU pressure triggered by GetOrOpen's Add, also +// under m.mu). It must not touch the cache. Marks the entry detached and +// closes the underlying pool iff no borrower remains. +func onCacheEvict[V io.Closer](_ PoolKey, e *entry[V]) { + e.evicted.Store(true) + e.maybeClose() +} + +// GetOrOpen returns the cached pool for (resource, method), constructing one +// via the factory on miss. The returned release MUST be called when the +// caller is finished with the pool; on the last release after eviction the +// pool's Close is invoked. release is idempotent. +// +// Each call refreshes the entry's idle TTL. +// +// Construction is serialized under m.mu, so the factory is invoked at most +// once per concurrent miss for the same key — fine while pools are cheap to +// build. +func (m *PoolCache[V]) GetOrOpen(resource, method string) (V, func(), error) { + var zero V + k := PoolKey{Resource: resource, Method: method} + + m.mu.Lock() + defer m.mu.Unlock() + + if e, ok := m.cache.Get(k); ok { + // Re-Add resets ExpiresAt without firing onCacheEvict (Add on an + // existing key updates in place; see expirable.LRU.Add). + m.cache.Add(k, e) + e.refs.Add(1) + return e.pool, m.releaseFn(e), nil + } + + v, err := m.factory(resource, method) + if err != nil { + return zero, nil, err + } + e := &entry[V]{key: k, pool: v} + e.refs.Store(1) + m.cache.Add(k, e) // may trigger LRU eviction of another entry + return e.pool, m.releaseFn(e), nil +} + +func (m *PoolCache[V]) releaseFn(e *entry[V]) func() { + var once sync.Once + return func() { + once.Do(func() { + e.refs.Add(-1) + e.maybeClose() + }) + } +} + +// Close synchronously closes every cached pool and returns the first error +// encountered. Callers should ensure no borrowed handles are still in use +// (e.g., the gRPC server has drained). Outstanding handles will not +// double-close: their release becomes a no-op for the underlying pool thanks +// to entry.closer. +func (m *PoolCache[V]) Close() error { + m.mu.Lock() + defer m.mu.Unlock() + var firstErr error + for _, k := range m.cache.Keys() { + e, ok := m.cache.Peek(k) + if !ok { + continue + } + e.evicted.Store(true) + if err := e.close(); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +}
diff --git a/bigtable/accelerator/resourcemanager/resourcemanager_test.go b/bigtable/accelerator/resourcemanager/resourcemanager_test.go new file mode 100644 index 0000000..10a94c4 --- /dev/null +++ b/bigtable/accelerator/resourcemanager/resourcemanager_test.go
@@ -0,0 +1,334 @@ +package resourcemanager + +import ( + "errors" + "sync/atomic" + "testing" + "time" +) + +type fakeClient struct { + name string + closeErr error + closed atomic.Int32 +} + +func (f *fakeClient) Close() error { + f.closed.Add(1) + return f.closeErr +} + +// Tests below always use the same method string so the cache key reduces to +// resource; eviction/LRU/TTL behavior is unchanged by the method dimension. +const testMethod = "ReadRow" + +func TestGetOrOpen_MissThenHit(t *testing.T) { + var calls atomic.Int32 + m := NewPoolCache[*fakeClient](4, func(resource, method string) (*fakeClient, error) { + calls.Add(1) + return &fakeClient{name: resource}, nil + }) + + v1, rel1, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + defer rel1() + v2, rel2, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + defer rel2() + if v1 != v2 { + t.Errorf("second GetOrOpen returned a different instance") + } + if got := calls.Load(); got != 1 { + t.Errorf("factory call count = %d; want 1", got) + } +} + +func TestGetOrOpen_KeyedByResourceAndMethod(t *testing.T) { + // Same resource, different methods => different pools (different factory calls). + var calls atomic.Int32 + m := NewPoolCache[*fakeClient](4, func(resource, method string) (*fakeClient, error) { + calls.Add(1) + return &fakeClient{name: resource + ":" + method}, nil + }) + + v1, r1, err := m.GetOrOpen("t1", "ReadRow") + if err != nil { + t.Fatalf("GetOrOpen(t1, ReadRow) err: %v", err) + } + defer r1() + v2, r2, err := m.GetOrOpen("t1", "MutateRow") + if err != nil { + t.Fatalf("GetOrOpen(t1, MutateRow) err: %v", err) + } + defer r2() + if v1 == v2 { + t.Errorf("expected distinct pools for distinct methods on same resource") + } + if got := calls.Load(); got != 2 { + t.Errorf("factory call count = %d; want 2", got) + } +} + +func TestGetOrOpen_FactoryError(t *testing.T) { + wantErr := errors.New("boom") + m := NewPoolCache[*fakeClient](4, func(string, string) (*fakeClient, error) { + return nil, wantErr + }) + + if _, _, err := m.GetOrOpen("t1", testMethod); !errors.Is(err, wantErr) { + t.Errorf("err = %v; want %v", err, wantErr) + } + // A failed construction must not be cached: a retry should re-invoke the factory. + var calls atomic.Int32 + m2 := NewPoolCache[*fakeClient](4, func(string, string) (*fakeClient, error) { + calls.Add(1) + return nil, wantErr + }) + _, _, _ = m2.GetOrOpen("t1", testMethod) + _, _, _ = m2.GetOrOpen("t1", testMethod) + if got := calls.Load(); got != 2 { + t.Errorf("factory called %d times on repeated failure; want 2", got) + } +} + +func TestGetOrOpen_LRUEvictionClosesEvicted(t *testing.T) { + clients := map[string]*fakeClient{} + m := NewPoolCache[*fakeClient](2, func(resource, method string) (*fakeClient, error) { + c := &fakeClient{name: resource} + clients[resource] = c + return c, nil + }) + + // Fill cache (capacity 2), then add a third entry to evict the LRU one (t1). + // Releasing immediately means refs==0 by the time eviction fires, so the + // LRU evictee should close in-line under onCacheEvict. + for _, name := range []string{"t1", "t2", "t3"} { + _, release, err := m.GetOrOpen(name, testMethod) + if err != nil { + t.Fatalf("GetOrOpen(%s) err: %v", name, err) + } + release() + } + + if got := clients["t1"].closed.Load(); got != 1 { + t.Errorf("evicted client t1 Close count = %d; want 1", got) + } + if got := clients["t2"].closed.Load(); got != 0 { + t.Errorf("retained client t2 Close count = %d; want 0", got) + } + if got := clients["t3"].closed.Load(); got != 0 { + t.Errorf("retained client t3 Close count = %d; want 0", got) + } +} + +func TestGetOrOpen_LRUOrderingFromHit(t *testing.T) { + clients := map[string]*fakeClient{} + m := NewPoolCache[*fakeClient](2, func(resource, method string) (*fakeClient, error) { + c := &fakeClient{name: resource} + clients[resource] = c + return c, nil + }) + + // Insert t1, t2; touch t1 so t2 becomes the LRU; insert t3 -> t2 evicts. + get := func(name string) { + _, release, err := m.GetOrOpen(name, testMethod) + if err != nil { + t.Fatalf("GetOrOpen(%s) err: %v", name, err) + } + release() + } + get("t1") + get("t2") + get("t1") + get("t3") + + if got := clients["t2"].closed.Load(); got != 1 { + t.Errorf("t2 Close count = %d; want 1 (should have been LRU-evicted)", got) + } + if got := clients["t1"].closed.Load(); got != 0 { + t.Errorf("t1 Close count = %d; want 0 (recently touched)", got) + } +} + +func TestClose_ClosesAllAndReturnsFirstError(t *testing.T) { + wantErr := errors.New("close-fail") + c1 := &fakeClient{name: "t1", closeErr: wantErr} + c2 := &fakeClient{name: "t2"} + clients := map[string]*fakeClient{"t1": c1, "t2": c2} + m := NewPoolCache[*fakeClient](4, func(resource, method string) (*fakeClient, error) { + return clients[resource], nil + }) + _, r1, _ := m.GetOrOpen("t1", testMethod) + _, r2, _ := m.GetOrOpen("t2", testMethod) + r1() + r2() + + err := m.Close() + if !errors.Is(err, wantErr) { + t.Errorf("Close err = %v; want %v", err, wantErr) + } + if got := c1.closed.Load(); got != 1 { + t.Errorf("t1 Close count = %d; want 1", got) + } + if got := c2.closed.Load(); got != 1 { + t.Errorf("t2 Close count = %d; want 1", got) + } +} + +func TestNewPoolCache_NonPositiveCapacityFallsBackToDefault(t *testing.T) { + m := NewPoolCache[*fakeClient](0, func(resource, method string) (*fakeClient, error) { + return &fakeClient{name: resource}, nil + }) + // Insert one more than the default to confirm an eviction occurs, which + // proves the cache is bounded (not, e.g., zero-sized or unbounded). + for i := 0; i < DefaultPoolCacheSize+1; i++ { + name := string(rune('a' + i)) + _, release, err := m.GetOrOpen(name, testMethod) + if err != nil { + t.Fatalf("GetOrOpen(%s) err: %v", name, err) + } + release() + } +} + +// --- Refcounting ---------------------------------------------------------- + +func TestGetOrOpen_EvictionDeferredWhileHandleHeld(t *testing.T) { + clients := map[string]*fakeClient{} + m := NewPoolCache[*fakeClient](1, func(resource, method string) (*fakeClient, error) { + c := &fakeClient{name: resource} + clients[resource] = c + return c, nil + }) + + // Borrow t1 and DON'T release; force LRU eviction of t1 by adding t2. + _, release1, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + _, release2, err := m.GetOrOpen("t2", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t2) err: %v", err) + } + defer release2() + + if got := clients["t1"].closed.Load(); got != 0 { + t.Fatalf("t1 closed while a handle was still outstanding (count=%d)", got) + } + + // Releasing the outstanding handle must close exactly once. + release1() + if got := clients["t1"].closed.Load(); got != 1 { + t.Errorf("t1 Close count after release = %d; want 1", got) + } + + // Idempotent: a second release on the same handle must not double-close. + release1() + if got := clients["t1"].closed.Load(); got != 1 { + t.Errorf("t1 Close count after second release = %d; want 1 (idempotent)", got) + } +} + +func TestGetOrOpen_TwoBorrowersCloseOnLastRelease(t *testing.T) { + c1 := &fakeClient{name: "t1"} + m := NewPoolCache[*fakeClient](1, func(string, string) (*fakeClient, error) { return c1, nil }) + + _, r1a, _ := m.GetOrOpen("t1", testMethod) + _, r1b, _ := m.GetOrOpen("t1", testMethod) + + // Evict t1 while two handles outstanding. + _, r2, _ := m.GetOrOpen("t2", testMethod) + defer r2() + if got := c1.closed.Load(); got != 0 { + t.Fatalf("closed while 2 handles outstanding (count=%d)", got) + } + + r1a() + if got := c1.closed.Load(); got != 0 { + t.Fatalf("closed with 1 handle still outstanding (count=%d)", got) + } + r1b() + if got := c1.closed.Load(); got != 1 { + t.Errorf("t1 Close count after final release = %d; want 1", got) + } +} + +// --- TTL ------------------------------------------------------------------ + +// The expirable LRU sweeps in (ttl / 100) ticks from a background goroutine, +// so TTL tests pick a short TTL and tolerate a small wait. Flaky-on-load by +// nature; bump the multiplier if CI is noisy. +const testTTL = 100 * time.Millisecond + +func waitForClose(t *testing.T, c *fakeClient, want int32) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if c.closed.Load() == want { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("client %q Close count = %d; want %d after wait", c.name, c.closed.Load(), want) +} + +func TestTTL_IdleEntryEvictedAndClosed(t *testing.T) { + c1 := &fakeClient{name: "t1"} + m := newPoolCacheWithTTL[*fakeClient](4, testTTL, func(string, string) (*fakeClient, error) { return c1, nil }) + + _, release, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + release() + + waitForClose(t, c1, 1) +} + +func TestTTL_SlidingRefreshOnGet(t *testing.T) { + c1 := &fakeClient{name: "t1"} + m := newPoolCacheWithTTL[*fakeClient](4, testTTL, func(string, string) (*fakeClient, error) { return c1, nil }) + + // Touch the entry repeatedly inside the TTL window; it must NOT expire. + deadline := time.Now().Add(3 * testTTL) + for time.Now().Before(deadline) { + _, release, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + release() + time.Sleep(testTTL / 4) + } + if got := c1.closed.Load(); got != 0 { + t.Errorf("client closed while still being used; close count = %d", got) + } + + // Stop touching it — must eventually expire and close. + waitForClose(t, c1, 1) +} + +func TestTTL_ExpiryWhileHeldDefersCloseUntilRelease(t *testing.T) { + c1 := &fakeClient{name: "t1"} + m := newPoolCacheWithTTL[*fakeClient](4, testTTL, func(string, string) (*fakeClient, error) { return c1, nil }) + + _, release, err := m.GetOrOpen("t1", testMethod) + if err != nil { + t.Fatalf("GetOrOpen(t1) err: %v", err) + } + + // Wait long enough for the background sweeper to fire — but the borrowed + // handle is still outstanding, so close must NOT happen yet. + time.Sleep(3 * testTTL) + if got := c1.closed.Load(); got != 0 { + t.Fatalf("close fired while handle outstanding; count = %d", got) + } + + release() + if got := c1.closed.Load(); got != 1 { + t.Errorf("close count after final release = %d; want 1", got) + } +}