feat(selector): invalidate token cache via DB notifications

[atharrva01]

Summary

Wire the cachedFetcher in the sherdlock selector to Postgres token DB notifications so any INSERT, UPDATE, or DELETE on the tokens table immediately marks the cache dirty. The next selector query then bypasses freshnessInterval and fetches from the DB rather than waiting for the time-based TTL to expire.

Surgical cache updates (patching individual cache buckets on INSERT/DELETE without a full DB scan) are intentionally deferred to a follow-up PR. That path has correctness edge cases (Ristretto eviction, notification idempotency, filter-field coverage) that deserve dedicated coverage before merging.

Changes

• token/services/storage/db/driver/token.go: TokenNotifier.Subscribe now returns (func() error, error). The cancel func silences only this subscription via an atomic.Bool; other subscribers on the same notifier are unaffected. TokenRecordReference extended with WalletID, Type, Quantity (already emitted by the Postgres trigger).
• token/services/storage/db/sql/postgres/tokens.go: Subscribe wraps the callback with an atomic.Bool gate and maps the trigger payload fields to TokenRecordReference.
• token/services/selector/sherdlock/fetcher.go: cachedFetcher subscribes to the notifier on construction; onTokenChange sets dirty=1 for all operations; Close() calls the cancel func. mixedFetcher.Close() delegates; Manager.Stop() calls it via io.Closer assertion.
• token/services/selector/sherdlock/manager.go: notifierDisabled config flag lets operators opt out of LISTEN/NOTIFY and fall back to pure time-based refresh.
• Tests: dirty-flag notification tests (Insert/Delete mark dirty, dirty cleared after update, dirty restored on DB error). DB notifier tests updated for new Subscribe signature.

Benchmark (3x, 4 CPUs, WSL2/Postgres)

Write rate	Notifier	Lazy
10 wps	4.9 ms/op	5.5 ms/op
100 wps	5.7 ms/op	4.8 ms/op
1000 wps	6.9 ms/op	6.4 ms/op

At 1000 wps (worst case) the notifier is ~8% slower than lazy because every fetch hits the DB -- the dirty flag is set continuously. The performance is within the same order of magnitude at all write rates. The dirty flag coalesces multiple notifications between queries so DB call rate is bounded by query rate, not write rate.

[atharrva01]

Hey @adecaro, I noticed the cachedFetcher had no way to react to DB writes in real time, so I wired driver.TokenNotifier into it. The cache now marks itself stale immediately on any insert or delete, rather than waiting out the freshness interval.

[adecaro]

Hi @atharrva01 , thanks for picking this up 🙏

So, in a running system, we can have tokens being continuously added and this will trigger a refresh of the cache that will prevent the selector any read. I see a possible substantial degradation of performance here. But what if instead we take the input passed to the listener and surgically update the cache with the new entry? We would only need to keep the lock to add the new token for a very short period. What do you think?

It might be good to check the benchmarks here token/services/selector/benchmark_test.go. Not sure, if they can already catch the scenario this PR is concerned with.

What do you think? 😄

[atharrva01]

Hey @adecaro, you're right, I missed that.

The tricky part is that TokenRecordReference only gives us {TxID, Index} and no WalletID or token Type, which are what we need to find the right cache bucket. So we'd still need a DB call either way. For inserts the idea would be to call QueryTokenDetails with just that one ID to get the wallet and type, then append to the right bucket. For deletes we'd need a small reverse index alongside the cache (TxID:Index to cache key) so we can find and remove it without scanning everything.

Still much cheaper than a full table scan under the write lock for high-write scenarios, so the direction makes sense to me.

One thing worth flagging, the cache value is currently an immutable iterators.Slice, so to do in-place appends and removes we'd need to switch to a mutable slice covered by the existing mu. Let me know if that's the way to go.

Also the existing benchmarks only cover read throughput on a static cache, they wouldn't catch the write-contention case you're describing. I can add a concurrent-write benchmark to make the improvement visible.

Does this sound like the right direction?

[adecaro]

Hi @atharrva01 , if more is needed we can modify what is put inside TokenRecordReference by updating the storage service. Please, have an assessment there. What do you think?

[atharrva01]

Hey @adecaro, had a look at the storage layer. The reference gets populated from a Postgres LISTEN/NOTIFY trigger and right now it only emits tx_id and idx in the payload. The notifier already supports extra columns through its PrimaryKey struct so extending it is pretty straightforward, we just pass owner_wallet_id, token_type and quantity as additional columns to NewTokenNotifier, update Subscribe to map them, and add those three fields to TokenRecordReference. That gives the callback everything it needs for a surgical update with no extra DB call at all. 4 files touched total. Happy to go ahead if this sounds right!

[adecaro]

HI @atharrva01 , yes, let's try that. There is a limit to the amount of data that Postgres can send. Let's see if we are in the boundaries 🤞

[adecaro]

@atharrva01 , if you have cycles, please, think also about a benchmark that we can run against Postgres under different loads to see how the selector react. The current benchmarks work against a fake DB. What do you think?

[atharrva01]

Hey @adecaro, good news, the infrastructure is already there. manager_test.go already spins up a real Postgres container via StartPostgres and wires up the real TokenNotifier. I can add a BenchmarkSelectorWithConcurrentWrites in that same file using startManagers, drive concurrent inserts at different rates alongside selections, and measure throughput. That should make the improvement visible end to end. I'll include this alongside the surgical update changes.

[adecaro]

Hi @atharrva01 , I was double checking and we have an 8,000 bytes limit on the payload Postgres would send back with the current approach. We need to be carefully. I was also checking the performance of this mechanism under high load. It is not design to be efficient under high load so each use case must evaluate if it makes sense to relay on this feature.

So, it might make sense also to have the possibility to disable completely the token notifier in case it starts being a bottleneck.

[atharrva01]

Hey @adecaro, good catches. On the 8KB limit, individual token fields should fit fine but I'll add a safeguard so the trigger fails loudly rather than silently truncating anything unexpected.

For the disable option, the fetcher already falls back to time-based polling when the notifier is unavailable, I'll just wire an explicit config flag into the selector config so operators can turn it off cleanly if it becomes a bottleneck under load.

[atharrva01]

Hey @adecaro, addressed all three points!

For the surgical update, I extended TokenRecordReference with WalletID, Type, and Quantity so the callback can patch the right cache bucket directly without any extra DB call. Inserts append, deletes filter by TxID+Index. The dirty flag fallback is still there for Updates and for stores that don't wire the notifier (like SQLite).

For the 8KB limit, the trigger now raises an explicit exception before calling pg_notify if the payload exceeds 8000 bytes, so any overflow fails loudly instead of delivering a malformed message.

For the disable option, added tokenNotifierDisabled: true to the selector config. When set, the fetcher skips the notifier entirely and falls back to the time-based freshness interval.

Also added BenchmarkSelectorWithConcurrentWrites in manager_test.go against a real Postgres container. It seeds tokens, drives inserts at ~10/s in the background, and measures selector throughput end to end.

Let me know if anything needs tweaking!

[adecaro]

Hi @atharrva01 , thanks for the effort and patience. I will review ASAP 🙏

[atharrva01]

hey @adecaro , a gentle ping on this pr , whenever you get time , thanks :)

[adecaro]

Hi @EvanYan1024 ,
Given that you have pushed recently changes related to the selector service, I would like to ask you some feedback on this PR, if you can 🙏
I want to make sure we introduce something that is helpful especially when running the stack at scale.

Many thanks

[EvanYan1024]

Thanks for the ping @adecaro. Sharing one observation, for what it's worth:

One concern about the benchmark added in manager_test.go:

ticker := time.NewTicker(100 * time.Millisecond)
// ...
_ = storeOne(fmt.Sprintf("write-tx-%d", n), 0, "0x64")

This is 10 writes/s. Issue #884 carries the performance label and explicitly targets selector throughput under continuous DB writes. A production CBDC-style load
typically sits in the high-hundreds to low-thousands TPS range, and each transfer writes 2–4 token rows. So 10 wps is roughly two orders of magnitude below the workload
this PR aims to improve, and at that rate the surgical-update path barely contends for f.mu at all — the benchmark would look healthy even if the implementation carried
a real lock-contention bug.

Would it be reasonable to parameterize the write rate and add a baseline comparison? Something like:

for _, rate := range []int{10, 100, 1000} {
    b.Run(fmt.Sprintf("write_rate_%dwps_notifier", rate), func(b *testing.B) {
        runBench(b, rate, /*notifierEnabled*/ true)
    })
    b.Run(fmt.Sprintf("write_rate_%dwps_lazy", rate), func(b *testing.B) {
        runBench(b, rate, /*notifierEnabled*/ false) // tokenNotifierDisabled
    })
}

That gives six datapoints across two dimensions:

1. Three write rates (10 / 100 / 1000 wps) to see where the cache notifier helps vs. where it might start to hurt — a trade-off you already flagged earlier in this
   thread.
2. Notifier vs. lazy baseline at each rate so the PR can demonstrate a concrete throughput win, or at worst confirm parity with the kill-switch on.

Without something like this, I'm not sure we can fully claim the PR closes the performance part of #884 — the wiring is in place, but it hasn't been shown to move the
needle on the workload the issue targets.

Happy to help look at the numbers once they're in. Thanks for the work on this!

[EvanYan1024]

Took another pass and ran into three more concerns worth flagging before merge:

1. Cache value is a shared mutable slice (race risk)

In onTokenChange, the Insert/Delete branches store via f.cache.Add(key, append(toks, newTok)). Meanwhile the read path:

toks, ok := f.cache.Get(tokenKey(walletID, currency))
f.mu.RUnlock()
if ok {
    return iterators.Slice(toks).NewPermutation(), nil
}

releases the read lock before the iterator is consumed. iterators.Slice (in fabric-smart-client/.../iterators/slice.go) doesn't copy — it holds the same backing array,
and so does NewPermutation. If append reuses the existing capacity on the next Insert, the writer will land on a slot a concurrent reader's iterator still aliases. go
test -race should catch this.

Could the surgical update path build a fresh slice before storing? Roughly:

case dbdriver.Insert:
    toks, _ := f.cache.Get(key)
    newToks := make([]*token2.UnspentTokenInWallet, len(toks)+1)
    copy(newToks, toks)
    newToks[len(toks)] = newTok
    f.cache.Add(key, newToks)

That keeps every cache value an immutable snapshot, and readers stay safe after RUnlock.

2. No unsubscribe path for the notifier

NewCachedFetcher calls notifier.Subscribe(f.onTokenChange), but cachedFetcher doesn't expose a Close(), and Manager.Stop() only cancels the cleaner goroutine.
TokenNotifier deliberately exports UnsubscribeAll() (driver/token.go:270), so release is presumably expected on the consumer side.

Since fetcherProvider constructs one fetcher per TMS (fetcher.go:84-100) — not a process-wide singleton — restarting / re-creating TMSes will accumulate LISTEN
connections and callback references. Probably worth adding a cachedFetcher.Close() that calls notifier.UnsubscribeAll() and wiring it through the manager/provider
lifecycle.

3. RAISE EXCEPTION in the trigger blocks the underlying write

The new payload guard in notifier.go raises when octet_length(output) > 8000. Since the trigger is AFTER ... FOR EACH ROW running in the same transaction as the original
INSERT, a RAISE EXCEPTION will abort the token write itself.

Loud-fail is the right instinct over silent truncation, but making a notification failure abort the business write feels like an availability regression. A safer
fallback might be: when the payload would exceed the limit, emit a minimal envelope (just TG_OP + primary keys) and let the subscriber fall back to the dirty-flag
full-refresh path. That keeps NOTIFY non-fatal while still avoiding malformed messages.

---

(Direction still feels right overall — this should help shrink the visibility window we see on the prepare side under load.)

[atharrva01]

hey @EvanYan1024, all four points addressed, also @adecaro please have a look whenever you get a chance, thanks

the insert race was a dumb one, switched to make+copy so readers that already grabbed the old slice don't share backing array with a concurrent writer. the delete path was fine since toks[:0:0] forces a fresh alloc on append anyway.

for the close/unsubscribe, added Close() on cachedFetcher -> UnsubscribeAll(), mixedFetcher delegates it, and Manager.Stop() calls it via io.Closer type assertion so i didn't have to touch the TokenFetcher interface or regenerate mocks.

the RAISE EXCEPTION thing was genuinely bad on my part, didn't think through that it runs in the same txn as the write. changed it to RAISE WARNING + RETURN NULL so an oversized payload just drops the notification and falls back to time-based refresh.

and yeah, restructured the benchmark to sweep 10/100/1000 wps with notifier on and off so the comparison is actually meaningful at production-ish load.

[atharrva01]

@adecaro , Here are the benchmark results from a local run (-benchtime=3x, 4 CPUs, WSL2/Postgres):

Notifier and lazy are within the same order of magnitude across all write rates, no meaningful throughput regression from the LISTEN/NOTIFY path. At 1000 wps the notifier is ~17% faster, likely because the cache stays warm from continuous notifications and avoids the time-based staleness check on each query.

[atharrva01]

Hey @adecaro, just wanted to say this PR ended up being a really good learning experience. Going through the review cycle helped me understand things I wouldn't have picked up just reading the code.

The surgical cache update was the most interesting part. I initially went with a simple dirty flag and full refresh, but patching the cache bucket directly forced me to understand how the mu lock interacts with readers after RUnlock. The shared backing array race was something I completely missed, I was appending onto the existing slice assuming cache.Add was enough isolation, but the iterator holds the same array and a concurrent writer can land on a slot the reader is still walking. The make+copy fix is obvious in hindsight.

The RAISE EXCEPTION in the trigger was a genuine blind spot. I didn't think through that the trigger runs inside the same transaction as the INSERT, so an exception would abort the token store write entirely. RAISE WARNING + RETURN NULL was the right call.

Two things I caught while doing a final self-review: Subscribe now returns a targeted cancel func instead of UnsubscribeAll() so other subscribers on the same notifier aren't affected, and I added tests covering the actual slice mutation paths in onTokenChange, all previous notifier tests used empty WalletID so they only hit the dirty-flag fallback.

Thanks for the patience through the back-and-forth, learned a lot from this one.

[adecaro]

LGTM

[EvanYan1024]

Thanks for the work here — wiring the selector cache to token DB notifications is definitely the right direction, and I like that the PR also keeps a time-based fallback.

I had one concern after reading through the cache update path: the surgical INSERT/DELETE optimization seems to rely on the notification payload being equivalent to the canonical selector query, but today it only carries tx_id/idx/wallet/type/quantity. The normal SpendableTokensIteratorBy path also filters on things like owner, is_deleted, spendable, and supported ledger formats, so I’m a bit worried the incremental path could admit or retain entries that a full DB refresh would not.

There are also a couple of cache-consistency edge cases that might be worth tightening before merge:

• INSERT notifications are not idempotent by token ID, so duplicate notifications or notifications racing with a full refresh could produce duplicate or stale entries.
• If a wallet/type bucket was evicted from Ristretto, an INSERT notification currently recreates it with only the new token, which could make the mixed fetcher skip the lazy DB fallback while hiding older tokens.
• The unsubscribe handle currently silences the callback, but it does not remove the wrapper from the underlying subscriber list or close the PG LISTEN lifecycle. Maybe that is intended to be owned by the TokenStore lifecycle, but it would be good to make that contract explicit.

Given that, I wonder if we could make this PR a more conservative first step: use the notifier only to mark the cache dirty on INSERT/DELETE/UPDATE, so the next selector query refreshes immediately. That should still address the original freshness issue without waiting for freshnessInterval, while keeping the surgical cache update as a follow-up once the payload/filtering/idempotency semantics are fully covered.

Happy to be corrected if I missed part of the storage contract here.

cc @adecaro

[adecaro]

Hi @EvanYan1024 , thanks for these additional thoughts. Let me address them:

Thanks for the work here — wiring the selector cache to token DB notifications is definitely the right direction, and I like that the PR also keeps a time-based fallback.

I had one concern after reading through the cache update path: the surgical INSERT/DELETE optimization seems to rely on the notification payload being equivalent to the canonical selector query, but today it only carries tx_id/idx/wallet/type/quantity. The normal SpendableTokensIteratorBy path also filters on things like owner, is_deleted, spendable, and supported ledger formats, so I’m a bit worried the incremental path could admit or retain entries that a full DB refresh would not.

Very good point. Let's clarify. To the best of my understanding:

• In case of an INSERT, we know that the default value given to spendable is true and StoreToken is assumed to use the default for it. It should also be safe to assume that a token just inserted has a format that can be managed by the current TMS instance.
• In case of a DELETE, we might want to remove the token from cache. The token can still be in cache but not cannot be locked because someone else locked it already.
• In case of an UPDATE, we need to be careful and check the other fields to make sure we match the behavior of SpendableTokensIteratorBy.

The other strategy is to perform in all cases an additional query to retrieve to the token information. This way, the trigger can just return the tx_id and index.

There are also a couple of cache-consistency edge cases that might be worth tightening before merge:

• INSERT notifications are not idempotent by token ID, so duplicate notifications or notifications racing with a full refresh could produce duplicate or stale entries.

Good point, better to check here.

• If a wallet/type bucket was evicted from Ristretto, an INSERT notification currently recreates it with only the new token, which could make the mixed fetcher skip the lazy DB fallback while hiding older tokens.

Very interesting. We need to make sure the cache gets refreshed or if no cache is there, we can drop the event altogether.

• The unsubscribe handle currently silences the callback, but it does not remove the wrapper from the underlying subscriber list or close the PG LISTEN lifecycle. Maybe that is intended to be owned by the TokenStore lifecycle, but it would be good to make that contract explicit.

Fair point, indeed.

Given that, I wonder if we could make this PR a more conservative first step: use the notifier only to mark the cache dirty on INSERT/DELETE/UPDATE, so the next selector query refreshes immediately. That should still address the original freshness issue without waiting for freshnessInterval, while keeping the surgical cache update as a follow-up once the payload/filtering/idempotency semantics are fully covered.

Don't we risk serious contention here if many tokens are inserted/updated/deleted? The cache will be invalidated continuously, no? Actually, updateCache already manages incremental updates of the cache, we should be fine there.
I just want to avoid continuous calls to SpendableTokensIteratorBy. I would be careful here. What do you think?

So, if the above issue can be avoided, why not splitting the PR in two. Sure. @atharrva01 , what do you think?

Happy to be corrected if I missed part of the storage contract here.

cc @adecaro

Super interesting, thanks. I'll remove the approve until we decide the next steps. Thanks both 🙏

[atharrva01]

hi @adecaro Agreed, splitting makes sense. The surgical update path has real edge cases (eviction, idempotency, filter mismatch) that deserve proper coverage before merging. I'll clean this PR down to just the dirty-flag subscription and open a follow-up for the surgical updates once those are addressed.

shall i proceed?

[adecaro]

Hi @atharrva01 , please, go ahead. Double check what happens during the benchmark when the notifier triggers events. We need to make sure the system doesn't get overwhelmed. Thanks much. This is taking more time but I think we are learning a lot 🙏

[atharrva01]

Thanks @adecaro, proceeding with the split now.

On the benchmark concern: with dirty-flag-only, the dirty flag coalesces all notifications that arrive between two fetch calls no matter how many INSERT/DELETE/UPDATE events fire, we only do one full DB refresh per selector query that finds dirty=1. So DB call rate is bounded by query rate, not write rate. The existing benchmark results (notifier ±same order of magnitude as lazy, ~17% faster at 1000 wps) already cover this path since they measure fetch throughput. I'll strip the surgical code and re-run to confirm the numbers hold.

[atharrva01]

@adecaro done, surgical update code stripped, PR is now dirty-flag only.

Benchmark results (3x, 4 CPUs, WSL2/Postgres):

Write rate	Notifier	Lazy
10 wps	4.9 ms/op	5.5 ms/op
100 wps	5.7 ms/op	4.8 ms/op
1000 wps	6.9 ms/op	6.4 ms/op

At 1000 wps (worst case) notifier is ~8% slower than lazy. The reason: at sustained high write rates the dirty flag is set on every notification, so every fetch hits the DB, essentially equivalent to the lazy path with a small notification overhead. The key property is that the dirty flag coalesces all notifications that arrive between two fetch calls into a single DB refresh, so DB call rate is bounded by query rate, not write rate. No overwhelming.

What is in this PR now:

• Subscribe to token DB notifications, mark cache dirty on any INSERT/UPDATE/DELETE
• Targeted cancel func (atomic.Bool gate per subscription, not UnsubscribeAll)
• Close()/Manager.Stop() wiring
• notifierDisabled flag for opt-out
• Dirty-flag tests + updated Subscribe signature tests

Follow-up PR will cover surgical cache updates (INSERT/DELETE patching without full DB scan) once the eviction, idempotency, and filter-coverage edge cases raised by @EvanYan1024 are properly addressed.

[adecaro]

Hi @atharrva01 , thanks much, I'll review ASAP 🙏

[adecaro]

what about setting the dirty flag only if a refresh is not already in progress? Could this improve the performance?

[adecaro]

Hi @atharrva01 , thanks for your effort and patience. I left a comment. I want to understand how to improve the performance by setting the dirty flag only in certain cases. What do you think? Thanks 🙏

[atharrva01]

hi @adecaro, good thought but there's a race hiding there, update() clears dirty before the DB call starts, so any token written during that window won't be in the snapshot. if we skip setting dirty when isUpdating is true, that token gets silently dropped and the cache stays stale until the next time-based refresh. the always-set behavior is intentional for exactly this reason. also checking isUpdating from the callback would need f.mu, which adds contention on every notification.

[AkramBitar]

Hello @adecaro,

I need your help with this PR. Could you please have look at this PR? What next?

Regards,
Akram