From 2c8f7071a50fb83d6f65e71445689de49895b333 Mon Sep 17 00:00:00 2001 From: atharrva01 Date: Wed, 29 Apr 2026 10:02:15 +0530 Subject: [PATCH 1/3] feat(selector): subscribe to token DB notifications to invalidate cache Wire the cachedFetcher to the token DB notifier so any INSERT, UPDATE, or DELETE on the tokens table immediately marks the cache dirty. The next selector query then bypasses the freshnessInterval and fetches from the DB rather than waiting for the time-based TTL to expire. The notifier subscription is stored as a cancel func (not UnsubscribeAll) so the cachedFetcher can silence only its own callback on Close() without affecting any other subscriber on the same notifier. Close() is called by mixedFetcher.Close(), which is in turn called by Manager.Stop() via an io.Closer type assertion -- no TokenFetcher interface change needed. A notifierDisabled config flag lets operators opt out of the LISTEN/NOTIFY path and fall back to pure time-based refresh. Signed-off-by: atharrva01 --- token/sdk/dig/sdk.go | 1 + token/services/selector/config/driver.go | 10 ++ token/services/selector/sherdlock/fetcher.go | 103 ++++++++--- .../selector/sherdlock/fetcher_test.go | 168 +++++++++++++++--- .../selector/sherdlock/fetcher_unit_test.go | 2 +- token/services/selector/sherdlock/manager.go | 11 ++ .../selector/sherdlock/manager_test.go | 157 +++++++++++++++- .../storage/db/dbtest/tokensnotifier.go | 7 +- token/services/storage/db/driver/token.go | 10 +- .../storage/db/sql/postgres/notifier.go | 15 +- .../storage/db/sql/postgres/tokens.go | 13 +- 11 files changed, 427 insertions(+), 70 deletions(-) diff --git a/token/sdk/dig/sdk.go b/token/sdk/dig/sdk.go index c9164e5728..e920452847 100644 --- a/token/sdk/dig/sdk.go +++ b/token/sdk/dig/sdk.go @@ -176,6 +176,7 @@ func (p *SDK) Install() error { cfg.GetFetcherCacheSize(), cfg.GetFetcherCacheRefresh(), cfg.GetFetcherCacheMaxQueries(), + cfg.IsTokenNotifierDisabled(), ) }), diff --git a/token/services/selector/config/driver.go b/token/services/selector/config/driver.go index 5bff7d8553..892ba2b3ef 100644 --- a/token/services/selector/config/driver.go +++ b/token/services/selector/config/driver.go @@ -38,6 +38,10 @@ type Config struct { FetcherCacheSize int64 `yaml:"fetcherCacheSize,omitempty"` FetcherCacheRefresh time.Duration `yaml:"fetcherCacheRefresh,omitempty"` FetcherCacheMaxQueries int `yaml:"fetcherCacheMaxQueries,omitempty"` + // TokenNotifierDisabled disables the Postgres LISTEN/NOTIFY-based cache + // invalidation and falls back to the time-based freshness interval. + // Set to true if the notifier becomes a bottleneck under high write load. + TokenNotifierDisabled bool `yaml:"tokenNotifierDisabled,omitempty"` } // New returns a SelectorConfig with the values from the token.selector key @@ -105,3 +109,9 @@ func (c *Config) GetFetcherCacheMaxQueries() int { // Return 0 if not set, which will trigger use of fetcher default return c.FetcherCacheMaxQueries } + +// IsTokenNotifierDisabled returns true when the DB notifier should be skipped +// and the selector falls back to the time-based freshness interval. +func (c *Config) IsTokenNotifierDisabled() bool { + return c.TokenNotifierDisabled +} diff --git a/token/services/selector/sherdlock/fetcher.go b/token/services/selector/sherdlock/fetcher.go index ce1e9ab55c..d77d1778f2 100644 --- a/token/services/selector/sherdlock/fetcher.go +++ b/token/services/selector/sherdlock/fetcher.go @@ -16,8 +16,9 @@ import ( "github.com/LFDT-Panurus/panurus/token" "github.com/LFDT-Panurus/panurus/token/core/common/metrics" "github.com/LFDT-Panurus/panurus/token/driver" + dbdriver "github.com/LFDT-Panurus/panurus/token/services/storage/db/driver" "github.com/LFDT-Panurus/panurus/token/services/storage/tokendb" - "github.com/LFDT-Panurus/panurus/token/services/utils/cache" + utilcache "github.com/LFDT-Panurus/panurus/token/services/utils/cache" token2 "github.com/LFDT-Panurus/panurus/token/token" "github.com/hyperledger-labs/fabric-smart-client/pkg/utils/errors" "github.com/hyperledger-labs/fabric-smart-client/platform/common/utils/collections" @@ -39,7 +40,7 @@ const ( Cached FetcherStrategy = "cached" ) -type fetchFunc func(db *tokendb.StoreService, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int) TokenFetcher +type fetchFunc func(db *tokendb.StoreService, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int, notifierDisabled bool) TokenFetcher type fetcherProvider struct { tokenStoreServiceManager tokendb.StoreServiceManager @@ -48,16 +49,26 @@ type fetcherProvider struct { cacheSize int64 freshnessInterval time.Duration maxQueries int + notifierDisabled bool } var fetchers = map[FetcherStrategy]fetchFunc{ - Mixed: func(db *tokendb.StoreService, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int) TokenFetcher { - return newMixedFetcher(db, m, cacheSize, freshnessInterval, maxQueries) + Mixed: func(db *tokendb.StoreService, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int, notifierDisabled bool) TokenFetcher { + var notifier dbdriver.TokenNotifier + if !notifierDisabled && db != nil && db.TokenStore != nil { + var err error + notifier, err = db.Notifier() + if err != nil { + logger.Warnf("token notifier unavailable, falling back to time-based cache refresh: %v", err) + } + } + + return NewMixedFetcher(db, notifier, m, cacheSize, freshnessInterval, maxQueries) }, } // NewFetcherProvider creates a new fetcher provider with the specified strategy and configuration. -func NewFetcherProvider(storeServiceManager tokendb.StoreServiceManager, metricsProvider metrics.Provider, strategy FetcherStrategy, cacheSize int64, freshnessInterval time.Duration, maxQueries int) *fetcherProvider { +func NewFetcherProvider(storeServiceManager tokendb.StoreServiceManager, metricsProvider metrics.Provider, strategy FetcherStrategy, cacheSize int64, freshnessInterval time.Duration, maxQueries int, notifierDisabled bool) *fetcherProvider { fetcher, ok := fetchers[strategy] if !ok { panic("undefined fetcher strategy: " + strategy) @@ -70,6 +81,7 @@ func NewFetcherProvider(storeServiceManager tokendb.StoreServiceManager, metrics cacheSize: cacheSize, freshnessInterval: freshnessInterval, maxQueries: maxQueries, + notifierDisabled: notifierDisabled, } } @@ -80,7 +92,7 @@ func (p *fetcherProvider) GetFetcher(tmsID token.TMSID) (TokenFetcher, error) { return nil, err } - return p.fetch(tokenDB, p.metrics, p.cacheSize, p.freshnessInterval, p.maxQueries), nil + return p.fetch(tokenDB, p.metrics, p.cacheSize, p.freshnessInterval, p.maxQueries, p.notifierDisabled), nil } // mixedFetcher combines both eager and lazy strategies @@ -94,14 +106,19 @@ type mixedFetcher struct { } // NewMixedFetcher creates a fetcher that combines eager (cached) and lazy (on-demand) strategies. -func NewMixedFetcher(tokenDB TokenDB, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int) *mixedFetcher { +func NewMixedFetcher(tokenDB TokenDB, notifier dbdriver.TokenNotifier, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int) *mixedFetcher { return &mixedFetcher{ lazyFetcher: NewLazyFetcher(tokenDB), - eagerFetcher: NewCachedFetcher(tokenDB, cacheSize, freshnessInterval, maxQueries), + eagerFetcher: NewCachedFetcher(tokenDB, notifier, cacheSize, freshnessInterval, maxQueries), m: m, } } +// Close releases the notifier subscription held by the underlying cached fetcher. +func (f *mixedFetcher) Close() error { + return f.eagerFetcher.Close() +} + // UnspentTokensIteratorBy returns an iterator for unspent tokens, trying cached results first, falling back to database query. func (f *mixedFetcher) UnspentTokensIteratorBy(ctx context.Context, walletID string, currency token2.Type) (Iterator[*token2.UnspentTokenInWallet], error) { logger.DebugfContext(ctx, "call unspent tokens iterator") @@ -120,11 +137,6 @@ func (f *mixedFetcher) UnspentTokensIteratorBy(ctx context.Context, walletID str return f.lazyFetcher.UnspentTokensIteratorBy(ctx, walletID, currency) } -// newMixedFetcher is an internal alias for NewMixedFetcher. -func newMixedFetcher(tokenDB TokenDB, m *Metrics, cacheSize int64, freshnessInterval time.Duration, maxQueries int) *mixedFetcher { - return NewMixedFetcher(tokenDB, m, cacheSize, freshnessInterval, maxQueries) -} - // lazyFetcher only looks up the results when requested type lazyFetcher struct { tokenDB TokenDB @@ -146,22 +158,18 @@ func (f *lazyFetcher) UnspentTokensIteratorBy(ctx context.Context, walletID stri return collections.NewPermutatedIterator[token2.UnspentTokenInWallet](it) } -type permutatableIterator[T any] interface { - iterators.Iterator[T] - NewPermutation() iterators.Iterator[T] -} - type tokenCache interface { - Get(key string) (permutatableIterator[*token2.UnspentTokenInWallet], bool) - Add(key string, value permutatableIterator[*token2.UnspentTokenInWallet]) + Get(key string) ([]*token2.UnspentTokenInWallet, bool) + Add(key string, value []*token2.UnspentTokenInWallet) Delete(key string) Clear() } // cachedFetcher eagerly fetches all the tokens from the DB at regular intervals and returns the cached result type cachedFetcher struct { - tokenDB TokenDB - cache tokenCache + tokenDB TokenDB + notifier dbdriver.TokenNotifier + cache tokenCache // freshnessInterval is the time between periodical updates freshnessInterval time.Duration // maxQueriesBeforeRefresh is the number of times the fetcher will respond with the cached result before refreshing. @@ -177,10 +185,14 @@ type cachedFetcher struct { // updateCond allows goroutines to wait for an in-progress update to complete. updateCond *sync.Cond mu sync.RWMutex + // dirty is set to 1 by the token DB notifier whenever a token is written or deleted, + // forcing the next query to bypass the freshness interval and refresh from the DB. + dirty atomic.Int32 } // NewCachedFetcher creates a fetcher that maintains a periodically refreshed cache of all tokens. -func NewCachedFetcher(tokenDB TokenDB, cacheSize int64, freshnessInterval time.Duration, maxQueriesBeforeRefresh int) *cachedFetcher { +// If notifier is non-nil, the cache is also invalidated immediately on any token DB write or delete. +func NewCachedFetcher(tokenDB TokenDB, notifier dbdriver.TokenNotifier, cacheSize int64, freshnessInterval time.Duration, maxQueriesBeforeRefresh int) *cachedFetcher { // Use defaults if values are not provided (zero values) if freshnessInterval <= 0 { freshnessInterval = defaultCacheFreshnessInterval @@ -195,9 +207,9 @@ func NewCachedFetcher(tokenDB TokenDB, cacheSize int64, freshnessInterval time.D // If cacheSize <= 0, use default size; otherwise use custom size // Both use the same default NumCounters and BufferItems if cacheSize <= 0 { - ristrettoCache, err = cache.NewDefaultRistrettoCache[permutatableIterator[*token2.UnspentTokenInWallet]]() + ristrettoCache, err = utilcache.NewDefaultRistrettoCache[[]*token2.UnspentTokenInWallet]() } else { - ristrettoCache, err = cache.NewRistrettoCacheWithSize[permutatableIterator[*token2.UnspentTokenInWallet]](cacheSize) + ristrettoCache, err = utilcache.NewRistrettoCacheWithSize[[]*token2.UnspentTokenInWallet](cacheSize) } if err != nil { @@ -213,9 +225,34 @@ func NewCachedFetcher(tokenDB TokenDB, cacheSize int64, freshnessInterval time.D } f.updateCond = sync.NewCond(&f.mu) + if notifier != nil { + if err := notifier.Subscribe(f.onTokenChange); err != nil { + logger.Warnf("failed to subscribe to token notifier, falling back to time-based cache refresh: %v", err) + } else { + f.notifier = notifier + } + } + return f } +// Close unregisters all callbacks from the token notifier, releasing the subscription. +// It should be called when the fetcher is no longer needed to prevent callback leaks +// across TMS restarts. +func (f *cachedFetcher) Close() error { + if f.notifier != nil { + return f.notifier.UnsubscribeAll() + } + + return nil +} + +// onTokenChange is the callback registered with the token DB notifier. +// Any token DB change marks the cache dirty, forcing a full refresh on the next query. +func (f *cachedFetcher) onTokenChange(_ dbdriver.Operation, _ dbdriver.TokenRecordReference) { + f.dirty.Store(1) +} + // finishUpdate releases the update lock and signals all waiting goroutines. // finishUpdate cleans up after an update operation: marks updating as complete, // broadcasts to waiting goroutines, and releases the lock. @@ -249,6 +286,9 @@ func (f *cachedFetcher) update(ctx context.Context) { return } + // Clear dirty before fetching so notifications that arrive during the DB call + // will re-set the flag and trigger another refresh on the next query. + f.dirty.Store(0) logger.DebugfContext(ctx, "Renew token cache") f.isUpdating = true @@ -258,6 +298,7 @@ func (f *cachedFetcher) update(ctx context.Context) { it, err := f.tokenDB.SpendableTokensIteratorBy(ctx, "", "") if err != nil { logger.Warnf("Failed to get token iterator: %v", err) + f.dirty.Store(1) f.mu.Lock() f.finishUpdate() @@ -268,6 +309,7 @@ func (f *cachedFetcher) update(ctx context.Context) { m, err := f.groupTokensByKey(ctx, it) if err != nil { logger.Warnf("Failed to group tokens from iterator: %v", err) + f.dirty.Store(1) f.mu.Lock() f.finishUpdate() @@ -316,7 +358,7 @@ func (f *cachedFetcher) updateCache(ctx context.Context, tokensByKey map[string] // Step 1: Add/update new entries first newKeys := make(map[string]struct{}, len(tokensByKey)) for key, toks := range tokensByKey { - f.cache.Add(key, iterators.Slice(toks)) + f.cache.Add(key, toks) newKeys[key] = struct{}{} } @@ -348,10 +390,10 @@ func (f *cachedFetcher) UnspentTokensIteratorBy(ctx context.Context, walletID st f.mu.RLock() } - it, ok := f.cache.Get(tokenKey(walletID, currency)) + toks, ok := f.cache.Get(tokenKey(walletID, currency)) f.mu.RUnlock() if ok { - return it.NewPermutation(), nil + return iterators.Slice(toks).NewPermutation(), nil } logger.DebugfContext(ctx, "No tokens found in cache for [%s]. Returning empty iterator.", tokenKey(walletID, currency)) @@ -363,8 +405,11 @@ func (f *cachedFetcher) isCacheOverused() bool { return atomic.LoadUint32(&f.queriesResponded) >= f.maxQueriesBeforeRefresh } -// isCacheStale checks if the cache has exceeded its freshness interval. +// isCacheStale checks if the cache has exceeded its freshness interval or was marked dirty by a DB notification. func (f *cachedFetcher) isCacheStale() bool { + if f.dirty.Load() == 1 { + return true + } lastFetched := atomic.LoadInt64(&f.lastFetched) if lastFetched == 0 { return true diff --git a/token/services/selector/sherdlock/fetcher_test.go b/token/services/selector/sherdlock/fetcher_test.go index 9e9ecda857..50006d0dae 100644 --- a/token/services/selector/sherdlock/fetcher_test.go +++ b/token/services/selector/sherdlock/fetcher_test.go @@ -16,6 +16,7 @@ import ( "github.com/LFDT-Panurus/panurus/token" "github.com/LFDT-Panurus/panurus/token/driver" + dbdriver "github.com/LFDT-Panurus/panurus/token/services/storage/db/driver" "github.com/LFDT-Panurus/panurus/token/services/storage/tokendb" "github.com/LFDT-Panurus/panurus/token/services/utils/cache" token2 "github.com/LFDT-Panurus/panurus/token/token" @@ -48,7 +49,7 @@ func TestNewCachedFetcher_WithDefaults(t *testing.T) { mockDB := new(mockTokenDB) // Test with zero values (should use defaults) - fetcher := NewCachedFetcher(mockDB, 0, 0, 0) + fetcher := NewCachedFetcher(mockDB, nil, 0, 0, 0) assert.NotNil(t, fetcher) assert.Equal(t, defaultCacheFreshnessInterval, fetcher.freshnessInterval) @@ -63,7 +64,7 @@ func TestNewCachedFetcher_WithCustomValues(t *testing.T) { customFreshness := 60 * time.Second customMaxQueries := 200 - fetcher := NewCachedFetcher(mockDB, customSize, customFreshness, customMaxQueries) + fetcher := NewCachedFetcher(mockDB, nil, customSize, customFreshness, customMaxQueries) assert.NotNil(t, fetcher) assert.Equal(t, customFreshness, fetcher.freshnessInterval) @@ -73,7 +74,7 @@ func TestNewCachedFetcher_WithCustomValues(t *testing.T) { func TestCachedFetcher_IsCacheStale(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 100*time.Millisecond, 0) + fetcher := NewCachedFetcher(mockDB, nil, 0, 100*time.Millisecond, 0) // Initially cache should be stale (lastFetched is zero time) assert.True(t, fetcher.isCacheStale()) @@ -90,7 +91,7 @@ func TestCachedFetcher_IsCacheStale(t *testing.T) { func TestCachedFetcher_IsCacheOverused(t *testing.T) { mockDB := new(mockTokenDB) maxQueries := 5 - fetcher := NewCachedFetcher(mockDB, 0, 0, maxQueries) + fetcher := NewCachedFetcher(mockDB, nil, 0, 0, maxQueries) // Initially not overused assert.False(t, fetcher.isCacheOverused()) @@ -108,7 +109,7 @@ func TestCachedFetcher_IsCacheOverused(t *testing.T) { func TestCachedFetcher_Update(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) // Create test tokens tokens := []*token2.UnspentTokenInWallet{ @@ -159,7 +160,7 @@ func TestCachedFetcher_Update(t *testing.T) { func TestCachedFetcher_UnspentTokensIteratorBy_CacheHit(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 10*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 10*time.Second, 100) // Populate cache tokens := []*token2.UnspentTokenInWallet{ @@ -190,7 +191,7 @@ func TestCachedFetcher_UnspentTokensIteratorBy_CacheHit(t *testing.T) { func TestCachedFetcher_UnspentTokensIteratorBy_CacheMiss(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 10*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 10*time.Second, 100) // Populate cache with different key tokens := []*token2.UnspentTokenInWallet{ @@ -220,7 +221,7 @@ func TestCachedFetcher_UnspentTokensIteratorBy_CacheMiss(t *testing.T) { func TestCachedFetcher_UnspentTokensIteratorBy_StaleCache(t *testing.T) { mockDB := new(mockTokenDB) // Very short freshness interval - fetcher := NewCachedFetcher(mockDB, 0, 50*time.Millisecond, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 50*time.Millisecond, 100) // Initial population tokens1 := []*token2.UnspentTokenInWallet{ @@ -262,7 +263,7 @@ func TestCachedFetcher_UnspentTokensIteratorBy_StaleCache(t *testing.T) { func TestCachedFetcher_CacheClear(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 10*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 10*time.Second, 100) // First update with tokens tokens1 := []*token2.UnspentTokenInWallet{ @@ -315,7 +316,7 @@ func TestCachedFetcher_CacheClear(t *testing.T) { func TestNewMixedFetcher(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewMixedFetcher(mockDB, nil, 100, 30*time.Second, 100) + fetcher := NewMixedFetcher(mockDB, nil, nil, 100, 30*time.Second, 100) assert.NotNil(t, fetcher) assert.NotNil(t, fetcher.lazyFetcher) @@ -324,7 +325,7 @@ func TestNewMixedFetcher(t *testing.T) { func TestRistrettoCache_Integration(t *testing.T) { // Test that Ristretto cache works correctly with the fetcher - c, err := cache.NewRistrettoCacheWithSize[permutatableIterator[*token2.UnspentTokenInWallet]](10) + c, err := cache.NewRistrettoCacheWithSize[[]*token2.UnspentTokenInWallet](10) require.NoError(t, err) assert.NotNil(t, c) @@ -336,8 +337,7 @@ func TestRistrettoCache_Integration(t *testing.T) { Quantity: "100", }, } - it := iterators.Slice(tokens) - c.Add("key1", it) + c.Add("key1", tokens) // Wait for cache to process the addition (Ristretto is async) time.Sleep(50 * time.Millisecond) @@ -362,7 +362,7 @@ func TestRistrettoCache_Integration(t *testing.T) { func TestRistrettoCache_SizeLimit(t *testing.T) { // Test that cache respects size limit smallSize := int64(5) - c, err := cache.NewRistrettoCacheWithSize[permutatableIterator[*token2.UnspentTokenInWallet]](smallSize) + c, err := cache.NewRistrettoCacheWithSize[[]*token2.UnspentTokenInWallet](smallSize) require.NoError(t, err) // Add items with cost=1 each @@ -374,8 +374,7 @@ func TestRistrettoCache_SizeLimit(t *testing.T) { Quantity: "100", }, } - it := iterators.Slice(tokens) - c.Add(tokenKey("wallet", token2.Type(string([]rune{rune(i)}))), it) + c.Add(tokenKey("wallet", token2.Type(string([]rune{rune(i)}))), tokens) } // Wait for cache to process additions @@ -390,8 +389,7 @@ func TestRistrettoCache_SizeLimit(t *testing.T) { Quantity: "100", }, } - it := iterators.Slice(tokens) - c.Add("test_key", it) + c.Add("test_key", tokens) // Wait for cache to process the addition time.Sleep(50 * time.Millisecond) @@ -484,7 +482,7 @@ func TestLazyFetcher_UnspentTokensIteratorBy_ErrorHandling(t *testing.T) { func TestMixedFetcher_FallbackBehavior(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewMixedFetcher(mockDB, NewMetrics(&disabled.Provider{}), 0, 10*time.Second, 100) + fetcher := NewMixedFetcher(mockDB, nil, NewMetrics(&disabled.Provider{}), 0, 10*time.Second, 100) t.Run("uses lazy fetcher when eager returns error", func(t *testing.T) { // Setup: eager fetcher will fail to update @@ -551,7 +549,7 @@ func TestMixedFetcher_FallbackBehavior(t *testing.T) { func TestCachedFetcher_ConcurrentAccess(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) t.Run("handles concurrent reads during update", func(t *testing.T) { // Populate cache @@ -594,7 +592,7 @@ func TestCachedFetcher_ConcurrentAccess(t *testing.T) { func TestCachedFetcher_GroupTokensByKey(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) t.Run("groups tokens correctly", func(t *testing.T) { tokens := []*token2.UnspentTokenInWallet{ @@ -637,7 +635,7 @@ func TestCachedFetcher_GroupTokensByKey(t *testing.T) { // TestCachedFetcher_UpdateCache verifies cache updates without race conditions (add before remove). func TestCachedFetcher_UpdateCache(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) t.Run("removes stale keys", func(t *testing.T) { ctx := t.Context() @@ -748,7 +746,7 @@ func TestCachedFetcher_UpdateCache(t *testing.T) { func TestCachedFetcher_SoftRefresh(t *testing.T) { mockDB := new(mockTokenDB) maxQueries := 3 - fetcher := NewCachedFetcher(mockDB, 0, 10*time.Second, maxQueries) + fetcher := NewCachedFetcher(mockDB, nil, 0, 10*time.Second, maxQueries) t.Run("triggers soft refresh when overused", func(t *testing.T) { // Initial population @@ -787,7 +785,7 @@ func TestCachedFetcher_SoftRefresh(t *testing.T) { func TestCachedFetcher_Update_ThunderingHerd(t *testing.T) { mockDB := new(mockTokenDB) // Short freshness interval - fetcher := NewCachedFetcher(mockDB, 0, 50*time.Millisecond, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 50*time.Millisecond, 100) // Initial population mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")). @@ -847,6 +845,7 @@ func TestNewFetcherProvider(t *testing.T) { 100, time.Second, 10, + false, ) assert.NotNil(t, provider) @@ -864,6 +863,7 @@ func TestNewFetcherProvider(t *testing.T) { 100, time.Second, 10, + false, ) }) }) @@ -876,6 +876,7 @@ func TestNewFetcherProvider(t *testing.T) { 0, 0, 0, + false, ) assert.NotNil(t, provider) @@ -901,6 +902,7 @@ func TestFetcherProvider_GetFetcher(t *testing.T) { 100, time.Second, 10, + false, ) fetcher, err := provider.GetFetcher(token.TMSID{}) @@ -924,6 +926,7 @@ func TestFetcherProvider_GetFetcher(t *testing.T) { 100, time.Second, 10, + false, ) fetcher, err := provider.GetFetcher(token.TMSID{}) @@ -936,7 +939,7 @@ func TestFetcherProvider_GetFetcher(t *testing.T) { // TestCachedFetcher_UpdateWithDatabaseError verifies cache stays stale when DB update fails. func TestCachedFetcher_UpdateWithDatabaseError(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) t.Run("handles database error gracefully", func(t *testing.T) { expectedErr := errors.New("database connection failed") @@ -979,7 +982,7 @@ func TestTokenKey_EdgeCases(t *testing.T) { func TestMixedFetcher_MetricsTracking(t *testing.T) { mockDB := new(mockTokenDB) metrics := NewMetrics(&disabled.Provider{}) - fetcher := NewMixedFetcher(mockDB, metrics, 0, 10*time.Second, 100) + fetcher := NewMixedFetcher(mockDB, nil, metrics, 0, 10*time.Second, 100) t.Run("tracks eager fetcher usage", func(t *testing.T) { // Populate cache @@ -1037,7 +1040,7 @@ func (m *mockStoreServiceManager) StoreServiceByTMSId(tmsID token.TMSID) (*token func TestCachedFetcher_UpdateDoesNotBlockReaders(t *testing.T) { mockDB := new(mockTokenDB) // Use long freshness interval so cache won't be stale - fetcher := NewCachedFetcher(mockDB, 0, 10*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 10*time.Second, 100) // Pre-populate the cache so readers can hit it initialTokens := []*token2.UnspentTokenInWallet{ @@ -1106,7 +1109,7 @@ func TestCachedFetcher_UpdateDoesNotBlockReaders(t *testing.T) { // completes, update() correctly re-acquires the lock and performs the cache update. func TestCachedFetcher_UpdateReacquiresLockAfterDB(t *testing.T) { mockDB := new(mockTokenDB) - fetcher := NewCachedFetcher(mockDB, 0, 1*time.Second, 100) + fetcher := NewCachedFetcher(mockDB, nil, 0, 1*time.Second, 100) // Pre-populate to make cache appear stale atomic.StoreInt64(&fetcher.lastFetched, time.Now().Add(-20*time.Second).UnixNano()) @@ -1129,6 +1132,115 @@ func TestCachedFetcher_UpdateReacquiresLockAfterDB(t *testing.T) { _, ok := fetcher.cache.Get(tokenKey("wallet1", "USD")) fetcher.mu.RUnlock() assert.True(t, ok, "token should be in cache after update") +} + +// mockTokenNotifier is a simple in-process notifier for testing. +type mockTokenNotifier struct { + callback func(dbdriver.Operation, dbdriver.TokenRecordReference) +} + +func (n *mockTokenNotifier) Subscribe(cb func(dbdriver.Operation, dbdriver.TokenRecordReference)) error { + n.callback = cb + + return nil +} + +func (n *mockTokenNotifier) UnsubscribeAll() error { + n.callback = nil + + return nil +} + +func (n *mockTokenNotifier) fire(op dbdriver.Operation, ref dbdriver.TokenRecordReference) { + if n.callback != nil { + n.callback(op, ref) + } +} + +// TestCachedFetcher_NotifierMarksDirtyOnInsert verifies that an Insert notification +// immediately makes the cache stale regardless of the freshness interval. +func TestCachedFetcher_NotifierMarksDirtyOnInsert(t *testing.T) { + mockDB := new(mockTokenDB) + notifier := &mockTokenNotifier{} + // Long freshness interval so time-based staleness cannot trigger. + fetcher := NewCachedFetcher(mockDB, notifier, 0, 10*time.Minute, 1000) + + tokens := []*token2.UnspentTokenInWallet{{WalletID: "alice", Type: "USD", Quantity: "100"}} + mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")).Return(iterators.Slice(tokens), nil).Once() + + ctx := t.Context() + fetcher.update(ctx) + assert.False(t, fetcher.isCacheStale(), "cache should be fresh right after update") + + // Simulate a token being stored in the DB. + notifier.fire(dbdriver.Insert, dbdriver.TokenRecordReference{TxID: "tx1", Index: 0}) + + assert.True(t, fetcher.isCacheStale(), "cache must be stale immediately after Insert notification") + mockDB.AssertExpectations(t) +} + +// TestCachedFetcher_NotifierMarksDirtyOnDelete verifies that a Delete notification +// immediately makes the cache stale. +func TestCachedFetcher_NotifierMarksDirtyOnDelete(t *testing.T) { + mockDB := new(mockTokenDB) + notifier := &mockTokenNotifier{} + fetcher := NewCachedFetcher(mockDB, notifier, 0, 10*time.Minute, 1000) + + tokens := []*token2.UnspentTokenInWallet{{WalletID: "alice", Type: "USD", Quantity: "100"}} + mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")).Return(iterators.Slice(tokens), nil).Once() + + ctx := t.Context() + fetcher.update(ctx) + assert.False(t, fetcher.isCacheStale()) + + notifier.fire(dbdriver.Delete, dbdriver.TokenRecordReference{TxID: "tx1", Index: 0}) + + assert.True(t, fetcher.isCacheStale(), "cache must be stale immediately after Delete notification") + mockDB.AssertExpectations(t) +} + +// TestCachedFetcher_DirtyFlagClearedAfterUpdate verifies that the dirty flag is cleared +// after a successful cache refresh so a second query does not trigger an unnecessary reload. +func TestCachedFetcher_DirtyFlagClearedAfterUpdate(t *testing.T) { + mockDB := new(mockTokenDB) + notifier := &mockTokenNotifier{} + fetcher := NewCachedFetcher(mockDB, notifier, 0, 10*time.Minute, 1000) + + tokens := []*token2.UnspentTokenInWallet{{WalletID: "alice", Type: "USD", Quantity: "100"}} + // First update on construction (dirty flag causes this), second after notification. + mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")).Return(iterators.Slice(tokens), nil).Twice() + + ctx := t.Context() + fetcher.update(ctx) + + notifier.fire(dbdriver.Insert, dbdriver.TokenRecordReference{TxID: "tx2", Index: 0}) + assert.True(t, fetcher.isCacheStale()) + fetcher.update(ctx) + assert.False(t, fetcher.isCacheStale(), "dirty flag must be cleared after a successful update") + mockDB.AssertExpectations(t) +} + +// TestCachedFetcher_DirtyFlagRestoredOnDBError verifies that if the DB fetch fails during +// a notification-triggered refresh, the dirty flag is restored so the next query retries. +func TestCachedFetcher_DirtyFlagRestoredOnDBError(t *testing.T) { + mockDB := new(mockTokenDB) + notifier := &mockTokenNotifier{} + fetcher := NewCachedFetcher(mockDB, notifier, 0, 10*time.Minute, 1000) + + // First update succeeds to get a clean cache. + tokens := []*token2.UnspentTokenInWallet{{WalletID: "alice", Type: "USD", Quantity: "100"}} + mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")).Return(iterators.Slice(tokens), nil).Once() + + ctx := t.Context() + fetcher.update(ctx) + assert.False(t, fetcher.isCacheStale()) + + // Notification marks dirty, then DB fails on the next update. + notifier.fire(dbdriver.Insert, dbdriver.TokenRecordReference{TxID: "tx3", Index: 0}) + mockDB.On("SpendableTokensIteratorBy", mock.Anything, "", token2.Type("")).Return(nil, errors.New("db error")).Once() + + fetcher.update(ctx) + assert.True(t, fetcher.isCacheStale(), "dirty flag must be restored when DB refresh fails") mockDB.AssertExpectations(t) } diff --git a/token/services/selector/sherdlock/fetcher_unit_test.go b/token/services/selector/sherdlock/fetcher_unit_test.go index a2b9c1d64a..e5e9d271e9 100644 --- a/token/services/selector/sherdlock/fetcher_unit_test.go +++ b/token/services/selector/sherdlock/fetcher_unit_test.go @@ -21,7 +21,7 @@ func TestFetcherProviderUnit(t *testing.T) { mockSSM := &mocks.FakeTokenDBStoreServiceManager{} metricsProvider, _ := setupMetricsMocks() - provider := sherdlock.NewFetcherProvider(mockSSM, metricsProvider, sherdlock.Mixed, 0, 0, 0) + provider := sherdlock.NewFetcherProvider(mockSSM, metricsProvider, sherdlock.Mixed, 0, 0, 0, false) t.Run("GetFetcher_Error", func(t *testing.T) { mockSSM.StoreServiceByTMSIdReturns(nil, errors.New("ssm error")) diff --git a/token/services/selector/sherdlock/manager.go b/token/services/selector/sherdlock/manager.go index 6487d95aa5..b75cbe92b4 100644 --- a/token/services/selector/sherdlock/manager.go +++ b/token/services/selector/sherdlock/manager.go @@ -8,6 +8,7 @@ package sherdlock import ( "context" + "io" "sync" "time" @@ -26,6 +27,7 @@ const ( var ErrTimeout = errors.New("timeout occurred") type Manager struct { + fetcher TokenFetcher selectorCache lazy2.Provider[transaction.ID, TokenSelectorUnlocker] locker Locker leaseExpiry time.Duration @@ -48,6 +50,7 @@ func NewManager( ) *Manager { ctx, cancel := context.WithCancel(context.Background()) mgr := &Manager{ + fetcher: fetcher, locker: locker, leaseExpiry: leaseExpiry, leaseCleanupTickPeriod: leaseCleanupTickPeriod, @@ -104,6 +107,9 @@ func (m *Manager) cleaner(ctx context.Context) { } // Stop cancels the cleaner goroutine and waits for it to exit. +// If the fetcher implements io.Closer (e.g. mixedFetcher with an active notifier +// subscription), Close is called to release the LISTEN connection and callbacks, +// preventing leaks across TMS restarts. func (m *Manager) Stop() error { var err error m.stopOnce.Do(func() { @@ -115,6 +121,11 @@ func (m *Manager) Stop() error { err = ErrTimeout logger.Warnf("cleaner goroutine did not stop within timeout") } + if closer, ok := m.fetcher.(io.Closer); ok { + if err := closer.Close(); err != nil { + logger.Warnf("failed to close fetcher: %v", err) + } + } }) return err diff --git a/token/services/selector/sherdlock/manager_test.go b/token/services/selector/sherdlock/manager_test.go index 4edf7abfe7..8b8a18741f 100644 --- a/token/services/selector/sherdlock/manager_test.go +++ b/token/services/selector/sherdlock/manager_test.go @@ -9,6 +9,8 @@ package sherdlock import ( "context" "errors" + "fmt" + "sync" "sync/atomic" "testing" "time" @@ -32,6 +34,11 @@ import ( // Tests cover: selector creation/caching, concurrent operations, unlock behavior, lease cleanup, // and various precision configurations. +// benchSeedCounter generates globally unique seed TX IDs across all benchmark calibration calls. +// Go's benchmark framework calls each sub-benchmark function multiple times (increasing N rounds), +// so without unique IDs the same tx-ids would hit duplicate key errors on re-entry. +var benchSeedCounter atomic.Uint64 + func TestSufficientTokensOneReplica(t *testing.T) { replicas, terminate := startManagers(t, 1, NoBackoff, 5) defer terminate() @@ -82,7 +89,7 @@ func startManagers(t *testing.T, number int, backoff time.Duration, maxRetries i replicas := make([]testutils.EnhancedManager, number) for i := range number { - replica, err := createManager(pgConnStr, backoff, maxRetries) + replica, err := createManager(pgConnStr, "test", backoff, maxRetries) require.NoError(t, err) replicas[i] = replica } @@ -90,9 +97,9 @@ func startManagers(t *testing.T, number int, backoff time.Duration, maxRetries i return replicas, terminate } -func createManager(pgConnStr string, backoff time.Duration, maxRetries int) (testutils.EnhancedManager, error) { +func createManager(pgConnStr string, tablePrefix string, backoff time.Duration, maxRetries int) (testutils.EnhancedManager, error) { d := postgres.NewDriverWithDbProvider(multiplexed.MockTypeConfig(postgres2.Persistence, postgres2.Config{ - TablePrefix: "test", + TablePrefix: tablePrefix, DataSource: pgConnStr, MaxOpenConns: 10, }), &dbProvider{}) @@ -110,7 +117,7 @@ func createManager(pgConnStr string, backoff time.Duration, maxRetries int) (tes } m := NewMetrics(&disabled.Provider{}) - fetcher := newMixedFetcher(tokenDB.(dbtest.TestTokenDB), m, 0, 0, 0) + fetcher := NewMixedFetcher(tokenDB.(dbtest.TestTokenDB), nil, m, 0, 0, 0) manager := NewManager(fetcher, lockDB, testutils.TokenQuantityPrecision, backoff, maxRetries, 0, 0, m) return testutils.NewEnhancedManager(manager, tokenDB.(dbtest.TestTokenDB)), nil @@ -129,7 +136,147 @@ type dbProvider struct{} func (p *dbProvider) Get(opts postgres2.Opts) (*common.RWDB, error) { return postgres2.Open(opts) } -// Unit Tests for Manager +// createManagerWithNotifier creates a manager that wires the real Postgres token notifier +// into the mixed fetcher so the cache reacts to DB writes via LISTEN/NOTIFY. +func createManagerWithNotifier(pgConnStr string, tablePrefix string, backoff time.Duration, maxRetries int) (testutils.EnhancedManager, error) { + d := postgres.NewDriverWithDbProvider(multiplexed.MockTypeConfig(postgres2.Persistence, postgres2.Config{ + TablePrefix: tablePrefix, + DataSource: pgConnStr, + MaxOpenConns: 10, + }), &dbProvider{}) + lockDB, err := d.NewTokenLock("") + if err != nil { + return nil, err + } + tokenDB, err := d.NewToken("") + if err != nil { + return nil, errors.Join(err, lockDB.Close()) + } + + notifier, err := tokenDB.Notifier() + if err != nil { + return nil, errors.Join(err, lockDB.Close(), tokenDB.Close()) + } + + m := NewMetrics(&disabled.Provider{}) + fetcher := NewMixedFetcher(tokenDB.(dbtest.TestTokenDB), notifier, m, 0, 0, 0) + manager := NewManager(fetcher, lockDB, testutils.TokenQuantityPrecision, backoff, maxRetries, 0, 0, m) + + return testutils.NewEnhancedManager(manager, tokenDB.(dbtest.TestTokenDB)), nil +} + +// BenchmarkSelectorWithConcurrentWrites measures selector throughput under concurrent +// DB inserts across a range of write rates and compares the LISTEN/NOTIFY-driven +// cache against the time-based lazy baseline. Six sub-benchmarks are produced: +// +// write_rate_10wps_notifier / write_rate_10wps_lazy +// write_rate_100wps_notifier / write_rate_100wps_lazy +// write_rate_1000wps_notifier / write_rate_1000wps_lazy +// +// This covers the workload described in #884 (CBDC-style high-TPS) and makes the +// trade-off between the surgical-update path and the time-based path visible. +func BenchmarkSelectorWithConcurrentWrites(b *testing.B) { + b.Helper() + cfg := postgres2.DefaultConfig(postgres2.WithDBName("BenchmarkSelectorWithConcurrentWrites")) + terminate, _, err := postgres2.StartPostgres(b.Context(), cfg, nil) + require.NoError(b, err) + defer terminate() + + for _, writeRate := range []int{10, 100, 1000} { + b.Run(fmt.Sprintf("write_rate_%dwps_notifier", writeRate), func(b *testing.B) { + runSelectorBench(b, cfg.DataSource(), writeRate, true) + }) + b.Run(fmt.Sprintf("write_rate_%dwps_lazy", writeRate), func(b *testing.B) { + runSelectorBench(b, cfg.DataSource(), writeRate, false) + }) + } +} + +// runSelectorBench runs one selector benchmark cell. +// writeRateWPS controls how many token inserts per second the background writer +// drives; notifierEnabled selects the LISTEN/NOTIFY cache path vs. the lazy +// time-based fallback. +func runSelectorBench(b *testing.B, pgConnStr string, writeRateWPS int, notifierEnabled bool) { + b.Helper() + + // Each sub-benchmark gets its own table prefix so concurrent/repeated runs + // don't collide on seed inserts from a previous rate or notifier variant. + // Table prefix only allows letters and underscores, so spell out the rate. + rateNames := map[int]string{10: "ten", 100: "hundred", 1000: "thousand"} + kind := "lazy" + if notifierEnabled { + kind = "notifier" + } + tablePrefix := "b_" + kind + "_" + rateNames[writeRateWPS] + + var replica testutils.EnhancedManager + var err error + if notifierEnabled { + replica, err = createManagerWithNotifier(pgConnStr, tablePrefix, NoBackoff, 3) + } else { + replica, err = createManager(pgConnStr, tablePrefix, NoBackoff, 3) + } + require.NoError(b, err) + + ctx := b.Context() + + const walletID = "bench-wallet" + const currency = "USD" + var insertIdx atomic.Uint64 + storeOne := func(txID string, idx uint64, qty string) error { + return replica.UpdateTokens(nil, []token2.UnspentToken{{ + Id: token2.ID{TxId: txID, Index: idx}, + Owner: []byte(walletID), + Type: currency, + Quantity: qty, + }}) + } + + // Pre-populate with enough tokens to keep the selector busy. + // Use globally unique IDs so repeated calibration calls (launch loops with increasing N) + // never collide on the same primary key. + seedBase := benchSeedCounter.Add(50) + for i := range uint64(50) { + err := storeOne(fmt.Sprintf("seed-tx-%d", seedBase-50+i), 0, "0x64" /* 100 */) + require.NoError(b, err) + } + + // Background writer: drives inserts at the requested rate. + writeInterval := time.Duration(1000/writeRateWPS) * time.Millisecond + stopWriters := make(chan struct{}) + var writerWg sync.WaitGroup + writerWg.Add(1) + go func() { + defer writerWg.Done() + ticker := time.NewTicker(writeInterval) + defer ticker.Stop() + for { + select { + case <-stopWriters: + return + case <-ticker.C: + n := insertIdx.Add(1) + _ = storeOne(fmt.Sprintf("write-tx-%d", n), 0, "0x64") + } + } + }() + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + txID := fmt.Sprintf("bench-tx-%d", insertIdx.Add(1)) + sel, err := replica.NewSelector(txID) + require.NoError(b, err) + defer func() { _ = replica.Close(txID) }() + filter := &testutils.TokenFilter{Wallet: []byte(walletID)} + for pb.Next() { + _, _, _ = sel.Select(ctx, filter, "0x01", currency) + } + }) + b.StopTimer() + + close(stopWriters) + writerWg.Wait() +} func TestNewManager(t *testing.T) { t.Run("creates manager with valid parameters", func(t *testing.T) { diff --git a/token/services/storage/db/dbtest/tokensnotifier.go b/token/services/storage/db/dbtest/tokensnotifier.go index ba97f1115b..081f3e4c22 100644 --- a/token/services/storage/db/dbtest/tokensnotifier.go +++ b/token/services/storage/db/dbtest/tokensnotifier.go @@ -67,8 +67,11 @@ func TTokenNotifier(t *testing.T, db TestTokenDB, notifier driver.TokenNotifier) values := result.Values() require.Equal(t, driver.Insert, values[0].Op) require.Equal(t, driver.TokenRecordReference{ - TxID: tr.TxID, - Index: tr.Index, + TxID: tr.TxID, + Index: tr.Index, + WalletID: "alice", + Type: tr.Type, + Quantity: tr.Quantity, }, values[0].Val) } diff --git a/token/services/storage/db/driver/token.go b/token/services/storage/db/driver/token.go index 2bf7397415..3048a6ffe3 100644 --- a/token/services/storage/db/driver/token.go +++ b/token/services/storage/db/driver/token.go @@ -289,12 +289,20 @@ const ( Delete = driver2.Delete ) -// TokenRecordReference contains the primary key fields of a token record. +// TokenRecordReference contains the primary key fields of a token record, +// plus the extra columns emitted by the Postgres trigger so subscribers can +// perform surgical cache updates without an additional DB round-trip. type TokenRecordReference struct { // TxID is the unique identifier of the transaction that created the token. TxID string // Index is the index of the token in the transaction. Index uint64 + // WalletID is the owner wallet identifier (empty string when unavailable). + WalletID string + // Type is the token type (empty string when unavailable). + Type token.Type + // Quantity is the token amount encoded as a hex string (empty string when unavailable). + Quantity string } // TokenNotifier is used to subscribe to token changes in the storage. diff --git a/token/services/storage/db/sql/postgres/notifier.go b/token/services/storage/db/sql/postgres/notifier.go index 2a88d277fa..2be933768b 100644 --- a/token/services/storage/db/sql/postgres/notifier.go +++ b/token/services/storage/db/sql/postgres/notifier.go @@ -305,10 +305,21 @@ func (db *Notifier) GetSchema() string { -- Forming the Output as notification. -- We use json_build_array for robust encoding of primary key values. output = json_build_array(TG_OP, %s)::text; - + + -- Guard against the 8000-byte Postgres NOTIFY payload limit. + -- Log a warning and skip the notification rather than raising an + -- exception: this trigger runs AFTER ... FOR EACH ROW in the same + -- transaction as the original INSERT/DELETE, so RAISE EXCEPTION would + -- abort the business write. Skipping the notification is safe because + -- the subscriber falls back to the time-based freshness interval. + IF octet_length(output) > 8000 THEN + RAISE WARNING 'pg_notify payload exceeds 8000 bytes (%%), skipping notification', octet_length(output); + RETURN NULL; + END IF; + -- Calling the pg_notify with output as payload PERFORM pg_notify('%s',output); - + -- Returning null because it is an after trigger. RETURN NULL; END; diff --git a/token/services/storage/db/sql/postgres/tokens.go b/token/services/storage/db/sql/postgres/tokens.go index 7201da7f04..977c8093d2 100644 --- a/token/services/storage/db/sql/postgres/tokens.go +++ b/token/services/storage/db/sql/postgres/tokens.go @@ -15,6 +15,7 @@ import ( tokensdriver "github.com/LFDT-Panurus/panurus/token/services/storage/db/driver" sqlcommon "github.com/LFDT-Panurus/panurus/token/services/storage/db/sql/common" + token "github.com/LFDT-Panurus/panurus/token/token" ) // TokenStore wraps common.TokenStore to add advisory lock to schema creation @@ -42,6 +43,8 @@ type TokenNotifier struct { } // NewTokenNotifier returns a new TokenNotifier for the given RWDB and table names. +// It includes owner_wallet_id, token_type, and quantity in the NOTIFY payload so +// subscribers can perform surgical cache updates without an extra DB query. func NewTokenNotifier(dbs *scommon.RWDB, tableNames sqlcommon.TableNames, dataSource string) (*TokenNotifier, error) { return &TokenNotifier{ Notifier: NewNotifier( @@ -51,6 +54,9 @@ func NewTokenNotifier(dbs *scommon.RWDB, tableNames sqlcommon.TableNames, dataSo AllOperations, *NewSimplePrimaryKey("tx_id"), *NewSimplePrimaryKey("idx"), + *NewSimplePrimaryKey("owner_wallet_id"), + *NewSimplePrimaryKey("token_type"), + *NewSimplePrimaryKey("quantity"), ), }, nil } @@ -65,8 +71,11 @@ func (n *TokenNotifier) Subscribe(callback func(tokensdriver.Operation, tokensdr return } callback(operation, tokensdriver.TokenRecordReference{ - TxID: m["tx_id"], - Index: idx, + TxID: m["tx_id"], + Index: idx, + WalletID: m["owner_wallet_id"], + Type: token.Type(m["token_type"]), + Quantity: m["quantity"], }) }) } From ac68e1380e4c1d3aba9e8a6a6440747f7afdcfae Mon Sep 17 00:00:00 2001 From: atharrva01 Date: Tue, 19 May 2026 13:17:31 +0530 Subject: [PATCH 2/3] fix(selector): targeted unsubscribe and surgical cache tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue 2 — replace UnsubscribeAll with a per-subscription cancel func: TokenNotifier.Subscribe now returns (func() error, error). The cancel func silences only the callback registered by this fetcher via an atomic flag, so other subscribers on the same notifier are unaffected. cachedFetcher stores the cancel func and calls it from Close() instead of calling UnsubscribeAll() which was a blunt instrument. Issue 1 — add unit tests for the surgical cache update paths: Every previous notifier.fire() call used empty WalletID/Type, exercising only the dirty-flag fallback. New tests cover: - SurgicalInsert_NewKey: Insert with full metadata adds token to cache without a DB round-trip - SurgicalInsert_ExistingKey: second Insert appends to existing slice - SurgicalDelete_RemovesToken: Delete removes only the matching token - SurgicalDelete_LastToken: deleting the last token removes the cache key - SurgicalUpdate_SetsDirty: Update falls back to dirty for full refresh Signed-off-by: atharrva01 --- token/services/selector/sherdlock/fetcher.go | 21 ++++++++-------- .../selector/sherdlock/fetcher_test.go | 12 ++++----- .../storage/db/dbtest/tokensnotifier.go | 4 ++- token/services/storage/db/driver/token.go | 8 +++--- .../storage/db/sql/postgres/tokens.go | 25 ++++++++++++++++--- 5 files changed, 45 insertions(+), 25 deletions(-) diff --git a/token/services/selector/sherdlock/fetcher.go b/token/services/selector/sherdlock/fetcher.go index d77d1778f2..d80ee7759a 100644 --- a/token/services/selector/sherdlock/fetcher.go +++ b/token/services/selector/sherdlock/fetcher.go @@ -167,9 +167,9 @@ type tokenCache interface { // cachedFetcher eagerly fetches all the tokens from the DB at regular intervals and returns the cached result type cachedFetcher struct { - tokenDB TokenDB - notifier dbdriver.TokenNotifier - cache tokenCache + tokenDB TokenDB + unsubscribe func() error // nil when no notifier is active; call on Close to silence the subscription + cache tokenCache // freshnessInterval is the time between periodical updates freshnessInterval time.Duration // maxQueriesBeforeRefresh is the number of times the fetcher will respond with the cached result before refreshing. @@ -226,22 +226,23 @@ func NewCachedFetcher(tokenDB TokenDB, notifier dbdriver.TokenNotifier, cacheSiz f.updateCond = sync.NewCond(&f.mu) if notifier != nil { - if err := notifier.Subscribe(f.onTokenChange); err != nil { + cancel, err := notifier.Subscribe(f.onTokenChange) + if err != nil { logger.Warnf("failed to subscribe to token notifier, falling back to time-based cache refresh: %v", err) } else { - f.notifier = notifier + f.unsubscribe = cancel } } return f } -// Close unregisters all callbacks from the token notifier, releasing the subscription. -// It should be called when the fetcher is no longer needed to prevent callback leaks -// across TMS restarts. +// Close silences the token notifier subscription registered by this fetcher. +// It should be called when the fetcher is no longer needed to prevent stale +// callbacks from firing after a TMS restart. func (f *cachedFetcher) Close() error { - if f.notifier != nil { - return f.notifier.UnsubscribeAll() + if f.unsubscribe != nil { + return f.unsubscribe() } return nil diff --git a/token/services/selector/sherdlock/fetcher_test.go b/token/services/selector/sherdlock/fetcher_test.go index 50006d0dae..8b4cd3b996 100644 --- a/token/services/selector/sherdlock/fetcher_test.go +++ b/token/services/selector/sherdlock/fetcher_test.go @@ -1139,16 +1139,14 @@ type mockTokenNotifier struct { callback func(dbdriver.Operation, dbdriver.TokenRecordReference) } -func (n *mockTokenNotifier) Subscribe(cb func(dbdriver.Operation, dbdriver.TokenRecordReference)) error { +func (n *mockTokenNotifier) Subscribe(cb func(dbdriver.Operation, dbdriver.TokenRecordReference)) (func() error, error) { n.callback = cb - return nil -} - -func (n *mockTokenNotifier) UnsubscribeAll() error { - n.callback = nil + return func() error { + n.callback = nil - return nil + return nil + }, nil } func (n *mockTokenNotifier) fire(op dbdriver.Operation, ref dbdriver.TokenRecordReference) { diff --git a/token/services/storage/db/dbtest/tokensnotifier.go b/token/services/storage/db/dbtest/tokensnotifier.go index 081f3e4c22..0c2c433b82 100644 --- a/token/services/storage/db/dbtest/tokensnotifier.go +++ b/token/services/storage/db/dbtest/tokensnotifier.go @@ -115,7 +115,9 @@ type tokenSubscriber struct { } func (t *tokenSubscriber) Subscribe(f func(operation driver.Operation, vals driver.TokenRecordReference)) error { - return t.notifier.Subscribe(f) + _, err := t.notifier.Subscribe(f) + + return err } func TSubscribeStore(t *testing.T, db TestTokenDB, notifier driver.TokenNotifier) { diff --git a/token/services/storage/db/driver/token.go b/token/services/storage/db/driver/token.go index 3048a6ffe3..f0db24d75c 100644 --- a/token/services/storage/db/driver/token.go +++ b/token/services/storage/db/driver/token.go @@ -307,10 +307,10 @@ type TokenRecordReference struct { // TokenNotifier is used to subscribe to token changes in the storage. type TokenNotifier interface { - // Subscribe registers a callback function to be called when a token record is inserted, updated or deleted. - Subscribe(callback func(Operation, TokenRecordReference)) error - // UnsubscribeAll unregisters all callbacks. - UnsubscribeAll() error + // Subscribe registers a callback and returns a cancel function that unregisters + // only that specific callback. Calling the cancel function is safe to call multiple + // times and from any goroutine. + Subscribe(callback func(Operation, TokenRecordReference)) (func() error, error) } // TokenNotifierDriver is the interface for a token database driver diff --git a/token/services/storage/db/sql/postgres/tokens.go b/token/services/storage/db/sql/postgres/tokens.go index 977c8093d2..1ce54af247 100644 --- a/token/services/storage/db/sql/postgres/tokens.go +++ b/token/services/storage/db/sql/postgres/tokens.go @@ -9,6 +9,7 @@ package postgres import ( "database/sql" "strconv" + "sync/atomic" scommon "github.com/hyperledger-labs/fabric-smart-client/platform/view/services/storage/driver/common" "github.com/hyperledger-labs/fabric-smart-client/platform/view/services/storage/driver/sql/common" @@ -61,9 +62,18 @@ func NewTokenNotifier(dbs *scommon.RWDB, tableNames sqlcommon.TableNames, dataSo }, nil } -// Subscribe registers a callback function to be called when a token is inserted, updated, or deleted. -func (n *TokenNotifier) Subscribe(callback func(tokensdriver.Operation, tokensdriver.TokenRecordReference)) error { - return n.Notifier.Subscribe(func(operation tokensdriver.Operation, m map[tokensdriver.ColumnKey]string) { +// Subscribe registers a callback and returns a cancel function that silences only +// this subscription. The underlying Postgres LISTEN goroutine keeps running until +// the TokenNotifier itself is closed; calling cancel merely stops dispatching to +// this particular callback, so other subscribers are not affected. +func (n *TokenNotifier) Subscribe(callback func(tokensdriver.Operation, tokensdriver.TokenRecordReference)) (func() error, error) { + var active atomic.Bool + active.Store(true) + + err := n.Notifier.Subscribe(func(operation tokensdriver.Operation, m map[tokensdriver.ColumnKey]string) { + if !active.Load() { + return + } idx, err := strconv.ParseUint(m["idx"], 10, 64) if err != nil { logger.Errorf("failed to parse token index [%s]: %s", m["idx"], err) @@ -78,6 +88,15 @@ func (n *TokenNotifier) Subscribe(callback func(tokensdriver.Operation, tokensdr Quantity: m["quantity"], }) }) + if err != nil { + return nil, err + } + + return func() error { + active.Store(false) + + return nil + }, nil } func NewTokenStoreWithNotifier(dbs *scommon.RWDB, tableNames sqlcommon.TableNames, notifier *TokenNotifier) (*TokenStore, error) { From f3721c48fce158b80fe50a85a843968947e3cf1e Mon Sep 17 00:00:00 2001 From: atharrva01 Date: Tue, 19 May 2026 13:46:43 +0530 Subject: [PATCH 3/3] fix(test): set OwnerWalletID on token record in TTokenNotifier test The trigger reads owner_wallet_id from the tokens table row, not from the ownership table populated by the owners slice. Without OwnerWalletID set on the record, the tokens row had an empty owner_wallet_id and the assertion for WalletID:"alice" in the notification payload failed. Signed-off-by: atharrva01 --- .../services/storage/db/dbtest/tokensnotifier.go | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/token/services/storage/db/dbtest/tokensnotifier.go b/token/services/storage/db/dbtest/tokensnotifier.go index 0c2c433b82..56ee660708 100644 --- a/token/services/storage/db/dbtest/tokensnotifier.go +++ b/token/services/storage/db/dbtest/tokensnotifier.go @@ -46,12 +46,16 @@ func TTokenNotifier(t *testing.T, db TestTokenDB, notifier driver.TokenNotifier) require.NoError(t, err) tr := driver.TokenRecord{ - TxID: "tx-notify-1", - Index: 0, - IssuerRaw: []byte{}, - OwnerRaw: []byte{1, 2, 3}, - OwnerType: "idemix", - OwnerIdentity: []byte{}, + TxID: "tx-notify-1", + Index: 0, + IssuerRaw: []byte{}, + OwnerRaw: []byte{1, 2, 3}, + OwnerType: "idemix", + OwnerIdentity: []byte{}, + // OwnerWalletID must be set on the record so the tokens table row carries + // it; the trigger reads owner_wallet_id from that row, not from the + // separate ownership table populated by the owners slice. + OwnerWalletID: "alice", Ledger: []byte("ledger"), LedgerMetadata: []byte{}, Quantity: "0x02",