From 5acd17a13007be8634e1ce7c78a376f1b80a4164 Mon Sep 17 00:00:00 2001 From: Jem Davies Date: Thu, 21 May 2026 16:35:28 +0100 Subject: [PATCH] impl lease & cache for buf schema registry Signed-off-by: Jem Davies --- internal/impl/protobuf/multimodule_watcher.go | 356 +++++++++++++++++- internal/impl/protobuf/processor_protobuf.go | 21 +- .../impl/protobuf/processor_protobuf_test.go | 67 +++- .../docs/components/processors/protobuf.md | 16 + 4 files changed, 445 insertions(+), 15 deletions(-) diff --git a/internal/impl/protobuf/multimodule_watcher.go b/internal/impl/protobuf/multimodule_watcher.go index 463ea7a269..5e7a906acf 100644 --- a/internal/impl/protobuf/multimodule_watcher.go +++ b/internal/impl/protobuf/multimodule_watcher.go @@ -1,16 +1,19 @@ package protobuf import ( + "bytes" "context" "errors" "fmt" "net/http" "strings" + "sync" "time" "buf.build/gen/go/bufbuild/reflect/connectrpc/go/buf/reflect/v1beta1/reflectv1beta1connect" connectrpc "connectrpc.com/connect" "github.com/bufbuild/prototransform" + "github.com/bufbuild/prototransform/leaser" "google.golang.org/protobuf/reflect/protoreflect" "google.golang.org/protobuf/reflect/protoregistry" @@ -23,7 +26,7 @@ type MultiModuleWatcher struct { var _ prototransform.Resolver = &MultiModuleWatcher{} -func newMultiModuleWatcher(bsrModules []*service.ParsedConfig) (*MultiModuleWatcher, error) { +func newMultiModuleWatcher(bsrModules []*service.ParsedConfig, mgr *service.Resources) (*MultiModuleWatcher, error) { if len(bsrModules) == 0 { return nil, errors.New("no modules provided") } @@ -53,7 +56,32 @@ func newMultiModuleWatcher(bsrModules []*service.ParsedConfig) (*MultiModuleWatc return nil, err } - watcher, err := newSchemaWatcher(context.Background(), bsrURL, bsrAPIKey, module, version) + var cacheName string + if cacheName, err = bsrModule.FieldString(fieldBsrCache); err != nil { + return nil, err + } + + var pollingPeriod time.Duration + if pollingPeriod, err = bsrModule.FieldDuration(fieldBsrPollingPeriod); err != nil { + return nil, err + } + + var leaser *bsrPollingLeaser + var cache *bsrCache + if cacheName != "" { + leaser = newLeaser(mgr, module, cacheName, + func() { + mgr.Logger().Infof("Acquired Lease for module: %v & cache: %v", module, cacheName) + }, + func() { + mgr.Logger().Infof("Released Lease for module: %v & cache: %v", module, cacheName) + }, + ) + + cache = newCache(mgr, cacheName) + } + + watcher, err := newSchemaWatcher(context.Background(), bsrURL, bsrAPIKey, module, version, cache, leaser, pollingPeriod) if err != nil { return nil, err } @@ -63,7 +91,7 @@ func newMultiModuleWatcher(bsrModules []*service.ParsedConfig) (*MultiModuleWatc return multiModuleWatcher, nil } -func newSchemaWatcher(ctx context.Context, bsrURL string, bsrAPIKey string, module string, version string) (*prototransform.SchemaWatcher, error) { +func newSchemaWatcher(ctx context.Context, bsrURL, bsrAPIKey, module, version string, cache *bsrCache, leaser *bsrPollingLeaser, pollingPeriod time.Duration) (*prototransform.SchemaWatcher, error) { // If no BSR url provided, extract from module if bsrURL == "" { segments := strings.Split(module, "/") @@ -88,8 +116,15 @@ func newSchemaWatcher(ctx context.Context, bsrURL string, bsrAPIKey string, modu module, version, ), - Jitter: 0.2, + Jitter: 0.2, + PollingPeriod: pollingPeriod, + } + + if cache != nil { + cfg.Leaser = leaser + cfg.Cache = cache } + watcher, err := prototransform.NewSchemaWatcher(ctx, cfg) if err != nil { return nil, fmt.Errorf("failed to create schema watcher: %w", err) @@ -173,3 +208,316 @@ func (w *MultiModuleWatcher) FindEnumByName(enum protoreflect.FullName) (protore } return nil, fmt.Errorf("could not find %s in any loaded modules", enum) } + +//------------------------------------------------------------------------------ + +var ( + cacheMap map[string]*bsrCache = make(map[string]*bsrCache) + cacheMapMu sync.Mutex +) + +// bsrCache implements prototransform.Cache +var _ prototransform.Cache = &bsrCache{} + +type bsrCache struct { + cacheName string + mgr *service.Resources +} + +func newCache(mgr *service.Resources, cacheName string) *bsrCache { + cacheMapMu.Lock() + defer cacheMapMu.Unlock() + if c, ok := cacheMap[cacheName]; ok { + return c + } + + c := &bsrCache{ + cacheName: cacheName, + mgr: mgr, + } + cacheMap[cacheName] = c + return c +} + +func (bc *bsrCache) Load(ctx context.Context, key string) ([]byte, error) { + var data []byte + var cacheErr error + err := bc.mgr.AccessCache(ctx, bc.cacheName, func(c service.Cache) { + data, cacheErr = c.Get(ctx, key) + }) + if err == nil { + err = cacheErr + } + if err != nil { + return nil, err + } + return data, nil +} + +func (bc *bsrCache) Save(ctx context.Context, key string, data []byte) error { + var cacheErr error + err := bc.mgr.AccessCache(ctx, bc.cacheName, func(c service.Cache) { + cacheErr = c.Set(ctx, key, data, nil) + }) + if err == nil { + err = cacheErr + } + return err +} + +//------------------------------------------------------------------------------ + +// leaseStore implements LeaseStore from the prototransform module +var _ leaser.LeaseStore = &leaseStore{} + +type leaseStore struct { + mgr *service.Resources + cacheName string +} + +func (ls *leaseStore) TryAcquire(ctx context.Context, leaseName string, holderId []byte, ttl time.Duration) (created bool, currentHolder []byte, err error) { + var cacheErr error + err = ls.mgr.AccessCache(ctx, ls.cacheName, func(c service.Cache) { + cacheErr = c.Add(ctx, leaseName, holderId, &ttl) + if cacheErr == nil { + created = true + currentHolder = holderId + return + } + if errors.Is(cacheErr, service.ErrKeyAlreadyExists) { + created = false + cacheErr = nil + return + } else if cacheErr != nil { + created = false + return + } + }) + if err == nil { + err = cacheErr + } + if err != nil { + return false, nil, err + } + + if !created { + err = ls.mgr.AccessCache(ctx, ls.cacheName, func(c service.Cache) { + currentHolder, cacheErr = c.Get(ctx, leaseName) + }) + if err == nil { + err = cacheErr + } + if err != nil { + return false, nil, err + } + } + + return created, currentHolder, err +} + +func (ls *leaseStore) Release(ctx context.Context, leaseName string, holderId []byte) error { + var cacheErr error + + err := ls.mgr.AccessCache(ctx, ls.cacheName, func(c service.Cache) { + var currentHolder []byte + currentHolder, cacheErr = c.Get(ctx, leaseName) + if errors.Is(cacheErr, service.ErrKeyNotFound) { + cacheErr = nil + return + } + if cacheErr != nil { + return + } + + if !bytes.Equal(currentHolder, holderId) { + return + } + + cacheErr = c.Delete(ctx, leaseName) + if errors.Is(cacheErr, service.ErrKeyNotFound) { + cacheErr = nil + } + }) + + if err == nil { + err = cacheErr + } + return err +} + +//------------------------------------------------------------------------------ + +// bsrPollingLeaser implements prototransform.Leaser +var _ prototransform.Leaser = &bsrPollingLeaser{} + +var ( + leaserMap map[string]*bsrPollingLeaser = make(map[string]*bsrPollingLeaser) + leaserMapMu sync.Mutex +) + +type bsrPollingLeaser struct { + leaseStore leaser.LeaseStore + leaseTTL time.Duration + pollingPeriod time.Duration + onAcquire func() + onRelease func() +} + +func newLeaser(mgr *service.Resources, module, cacheName string, onAcquire, onRelease func()) *bsrPollingLeaser { + leaserMapMu.Lock() + defer leaserMapMu.Unlock() + if l, ok := leaserMap[module+"::"+cacheName]; ok { + return l + } + + l := &bsrPollingLeaser{ + leaseStore: &leaseStore{ + mgr: mgr, + cacheName: cacheName, + }, + leaseTTL: time.Minute * 5, + pollingPeriod: time.Minute, + onAcquire: onAcquire, + onRelease: onRelease, + } + leaserMap[module+"::"+cacheName] = l + return l +} + +func (bpl *bsrPollingLeaser) NewLease(ctx context.Context, leaseName string, leaseHolder []byte) prototransform.Lease { + ctx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + newLease := &lease{ + cancel: cancel, + done: done, + err: prototransform.ErrLeaseStateNotYetKnown, + onAcquire: bpl.onAcquire, + onRelease: bpl.onRelease, + } + go newLease.run(ctx, bpl, leaseName, leaseHolder, done) + return newLease +} + +//------------------------------------------------------------------------------ +// This lease implementation has been taken from: +// https://github.com/bufbuild/prototransform/blob/main/leaser/polling_leaser.go +//------------------------------------------------------------------------------ + +// lease implements prototransform.Lease +var _ prototransform.Lease = &lease{} + +type lease struct { + cancel context.CancelFunc + done <-chan struct{} + + mu sync.Mutex + isHeld bool + err error + + notifyMu sync.Mutex + onAcquire, onRelease func() +} + +func (l *lease) IsHeld() (bool, error) { + l.mu.Lock() + isHeld, err := l.isHeld, l.err + l.mu.Unlock() + return isHeld, err +} + +func (l *lease) SetCallbacks(onAcquire, onRelease func()) { + l.mu.Lock() + defer l.mu.Unlock() + l.onAcquire, l.onRelease = onAcquire, onRelease + if l.isHeld && l.onAcquire != nil { + go func() { + l.notifyMu.Lock() + defer l.notifyMu.Unlock() + l.onAcquire() + }() + } +} + +func (l *lease) Cancel() { + l.cancel() + <-l.done +} + +func (l *lease) run(ctx context.Context, leaser *bsrPollingLeaser, key string, value []byte, done chan<- struct{}) { + defer close(done) + ticker := time.NewTicker(leaser.pollingPeriod) + defer ticker.Stop() + l.poll(ctx, leaser, key, value) + for { + select { + case <-ctx.Done(): + l.releaseNow(ctx, leaser, key, value) + return + case <-ticker.C: + if ctx.Err() != nil { + // skip polling if context is done + l.releaseNow(ctx, leaser, key, value) + return + } + l.poll(ctx, leaser, key, value) + } + } +} + +func (l *lease) poll(ctx context.Context, leaser *bsrPollingLeaser, key string, value []byte) { + created, holder, err := leaser.leaseStore.TryAcquire(ctx, key, value, leaser.leaseTTL) + if err != nil { + l.released(err) + return + } + if created { + l.acquired() + return + } + if bytes.Equal(holder, value) { + // The existing lease is ours + l.acquired() + return + } + // The existing lease is not ours + l.released(nil) +} + +func (l *lease) releaseNow(ctx context.Context, leaser *bsrPollingLeaser, key string, value []byte) { + l.mu.Lock() + isHeld := l.isHeld + l.mu.Unlock() + if isHeld { + // best effort: immediately release if we hold it + _ = leaser.leaseStore.Release(ctx, key, value) + } + l.released(nil) +} + +func (l *lease) acquired() { + l.mu.Lock() + defer l.mu.Unlock() + if !l.isHeld && l.onAcquire != nil { + go func() { + l.notifyMu.Lock() + defer l.notifyMu.Unlock() + l.onAcquire() + }() + } + l.isHeld = true + l.err = nil +} + +func (l *lease) released(err error) { + l.mu.Lock() + defer l.mu.Unlock() + if l.isHeld && l.onRelease != nil { + go func() { + l.notifyMu.Lock() + defer l.notifyMu.Unlock() + l.onRelease() + }() + } + l.isHeld = false + l.err = err +} diff --git a/internal/impl/protobuf/processor_protobuf.go b/internal/impl/protobuf/processor_protobuf.go index 59f9575754..dce41988e7 100644 --- a/internal/impl/protobuf/processor_protobuf.go +++ b/internal/impl/protobuf/processor_protobuf.go @@ -25,11 +25,13 @@ const ( fieldUseProtoNames = "use_proto_names" // BSR Config - fieldBsrConfig = "bsr" - fieldBsrModule = "module" - fieldBSRUrl = "url" - fieldBsrAPIKey = "api_key" - fieldBsrVersion = "version" + fieldBsrConfig = "bsr" + fieldBsrModule = "module" + fieldBSRUrl = "url" + fieldBsrAPIKey = "api_key" + fieldBsrVersion = "version" + fieldBsrCache = "cache" + fieldBsrPollingPeriod = "polling_period" ) func protobufProcessorSpec() *service.ConfigSpec { @@ -76,9 +78,16 @@ Attempts to create a target protobuf message from a generic JSON structure. Description("Buf Schema Registry server API key, can be left blank for a public registry."). Secret(). Default(""), + service.NewDurationField(fieldBsrPollingPeriod). + Description("The period in which to poll the Buf Schema Registry."). + Default("5m"). + Advanced(), service.NewStringField(fieldBsrVersion). Description("Version to retrieve from the Buf Schema Registry, leave blank for latest."). Default("").Advanced(), + service.NewStringField(fieldBsrCache). + Description("If set schemas retrieved from the BSR will be saved in a [cache resource](/docs/components/caches/about)"). + Default("").Advanced(), ).Description("Buf Schema Registry configuration. Either this field or `import_paths` must be populated. Note that this field is an array, and multiple BSR configurations can be provided."). Default([]any{}), ).LintRule(` @@ -508,7 +517,7 @@ func newProtobuf(conf *service.ParsedConfig, mgr *service.Resources) (*protobufP // if BSR config is present, use BSR to discover proto definitions if len(bsrModules) > 0 { - p.multiModuleWatcher, err = newMultiModuleWatcher(bsrModules) + p.multiModuleWatcher, err = newMultiModuleWatcher(bsrModules, mgr) if err != nil { return nil, fmt.Errorf("failed to create MultiModuleWatcher: %w", err) } diff --git a/internal/impl/protobuf/processor_protobuf_test.go b/internal/impl/protobuf/processor_protobuf_test.go index 28a3a8e997..e4278e34ff 100644 --- a/internal/impl/protobuf/processor_protobuf_test.go +++ b/internal/impl/protobuf/processor_protobuf_test.go @@ -7,6 +7,7 @@ import ( "net" "net/http" "strings" + "sync" "testing" "buf.build/gen/go/bufbuild/reflect/connectrpc/go/buf/reflect/v1beta1/reflectv1beta1connect" @@ -21,6 +22,9 @@ import ( "google.golang.org/protobuf/reflect/protoregistry" "google.golang.org/protobuf/types/descriptorpb" + "github.com/warpstreamlabs/bento/internal/component/testutil" + "github.com/warpstreamlabs/bento/internal/manager/mock" + "github.com/warpstreamlabs/bento/internal/message" "github.com/warpstreamlabs/bento/public/service" ) @@ -108,7 +112,7 @@ discard_unknown: %t }) t.Run(test.name+" bsr", func(t *testing.T) { - mockBSRServerAddress := runMockBSRServer(t) + mockBSRServerAddress, _ := runMockBSRServer(t) conf, err := protobufProcessorSpec().ParseYAML(fmt.Sprintf(` operator: from_json @@ -229,7 +233,7 @@ use_proto_names: %t }) t.Run(test.name+" bsr", func(t *testing.T) { - mockBSRServerAddress := runMockBSRServer(t) + mockBSRServerAddress, _ := runMockBSRServer(t) conf, err := protobufProcessorSpec().ParseYAML(fmt.Sprintf(` operator: to_json @@ -309,7 +313,7 @@ import_paths: [ %v ] }) t.Run(test.name+" bsr", func(tt *testing.T) { - mockBSRServerAddress := runMockBSRServer(t) + mockBSRServerAddress, _ := runMockBSRServer(t) conf, err := protobufProcessorSpec().ParseYAML(fmt.Sprintf(` operator: %v @@ -395,16 +399,69 @@ protobuf: } } +func TestBufCache(t *testing.T) { + mockBSRServerAddress, s := runMockBSRServer(t) + + mgr := mock.NewManager() + mgr.Caches["foo_cache"] = map[string]mock.CacheItem{} + + input := []byte{ + 0x8, 0xeb, 0x5, 0x12, 0x2a, 0xa, 0x21, 0x74, 0x79, 0x70, 0x65, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, + 0x65, 0x61, 0x70, 0x69, 0x73, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x74, 0x65, 0x73, 0x74, 0x69, 0x6e, + 0x67, 0x2e, 0x48, 0x6f, 0x75, 0x73, 0x65, 0x12, 0x5, 0x12, 0x3, 0x31, 0x32, 0x33, + } + + conf, err := testutil.ProcessorFromYAML(fmt.Sprintf(` +protobuf: + operator: to_json + message: testing.Envelope + bsr: + - module: testing + url: %s + cache: foo_cache + polling_period: 20s +`, "http://"+mockBSRServerAddress)) + require.NoError(t, err) + + for range 100 { + p, err := mgr.NewProcessor(conf) + require.NoError(t, err) + + msg := message.QuickBatch([][]byte{ + input, + }) + + msgs, res := p.ProcessBatch(context.Background(), msg) + require.NoError(t, res) + require.Len(t, msgs, 1) + + mBytes := msgs[0].Get(0).AsBytes() + + require.JSONEq(t, `{"id":747,"content":{"@type":"type.googleapis.com/testing.House","address":"123"}}`, string(mBytes)) + } + + s.ccMutex.Lock() + defer s.ccMutex.Unlock() + assert.LessOrEqual(t, s.cc, 5) +} + type fileDescriptorSetServer struct { fileDescriptorSet *descriptorpb.FileDescriptorSet + + cc int + ccMutex sync.Mutex } func (s *fileDescriptorSetServer) GetFileDescriptorSet(_ context.Context, request *connect.Request[v1beta1.GetFileDescriptorSetRequest]) (*connect.Response[v1beta1.GetFileDescriptorSetResponse], error) { + s.ccMutex.Lock() + s.cc++ + s.ccMutex.Unlock() + response := &v1beta1.GetFileDescriptorSetResponse{FileDescriptorSet: s.fileDescriptorSet, Version: request.Msg.GetVersion()} return connect.NewResponse(response), nil } -func runMockBSRServer(t *testing.T) string { +func runMockBSRServer(t *testing.T) (string, *fileDescriptorSetServer) { // load files into protoregistry.Files mockResources := service.MockResources() files, _, err := loadDescriptors(mockResources.FS(), []string{protosPath}) @@ -444,5 +501,5 @@ func runMockBSRServer(t *testing.T) string { } }() - return listener.Addr().String() + return listener.Addr().String(), fileDescriptorSetServer } diff --git a/website/docs/components/processors/protobuf.md b/website/docs/components/processors/protobuf.md index 048781478a..812b446ed8 100644 --- a/website/docs/components/processors/protobuf.md +++ b/website/docs/components/processors/protobuf.md @@ -315,6 +315,14 @@ This field contains sensitive information that usually shouldn't be added to a c Type: `string` Default: `""` +### `bsr[].polling_period` + +The period in which to poll the Buf Schema Registry. + + +Type: `string` +Default: `"5m"` + ### `bsr[].version` Version to retrieve from the Buf Schema Registry, leave blank for latest. @@ -323,4 +331,12 @@ Version to retrieve from the Buf Schema Registry, leave blank for latest. Type: `string` Default: `""` +### `bsr[].cache` + +If set schemas retrieved from the BSR will be saved in a [cache resource](/docs/components/caches/about) + + +Type: `string` +Default: `""` +