Skip to content

Commit f3c161b

Browse files
committed
fix(selector): close the token iterator displaced on retry
Selector.selectInternal replaced s.cache on every immediate retry without closing the iterator it displaced, so a single Select could abandon up to maxImmediateRetries iterators. On the lazy fetcher path an iterator wraps sql.Rows, so each one held a database cursor and its pooled connection until the finalizer ran — draining, under contention, the same pool the vault and commit path use. s.mu is documented as protecting the cache field, but selectInternal read s.cache.Next() and wrote s.cache without holding it. A concurrent Close() nils the field, so it could be dereferenced after the isClosed() check. Route both accesses through mutex-guarded helpers: next() reads the current iterator under s.mu, and swapCache() closes the displaced iterator before installing the new one. If the selector was closed while the fetch was in flight, swapCache closes the new iterator too and reports the error, so no iterator is left unclosed on any path. Fixes #2019 Signed-off-by: AkramBitar <akram@il.ibm.com>
1 parent 0dd324a commit f3c161b

3 files changed

Lines changed: 166 additions & 2 deletions

File tree

docs/services/selector.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,14 @@ The selector uses a **Token Fetcher** to retrieve available tokens from the data
163163
5. Selector iterates through tokens, attempting to lock each one
164164
6. If insufficient tokens, selector requests fresh data and retries
165165

166+
**Iterator lifecycle:** the iterator the fetcher hands out owns a resource — on the lazy
167+
path it wraps the query's `sql.Rows`, and therefore a database cursor and its pooled
168+
connection — so exactly one `Close()` per iterator is required. The selector holds at most
169+
one iterator at a time: step 6 above closes the iterator it displaces before installing the
170+
refreshed one, and `Selector.Close()` closes whichever is current. Both happen under the
171+
selector's mutex, so closing a selector while a retry is in flight neither leaks an iterator
172+
nor races with the retry.
173+
166174
**Adaptive refresh strategy** with two triggers:
167175
- **Time-based**: Refreshes when data is older than `fetcherCacheRefresh`
168176
- **Query-based**: Refreshes after `fetcherCacheMaxQueries` queries to prevent serving stale data in high-throughput scenarios

token/services/selector/sherdlock/selector.go

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ func (s *Selector) selectInternal(ctx context.Context, owner token.OwnerFilter,
164164
}
165165
sum, selected, tokensLockedByOthersExist, immediateRetries := token2.NewZeroQuantity(s.precision), collections.NewSet[*token2.ID](), true, 0
166166
for {
167-
if t, err := s.cache.Next(); err != nil {
167+
if t, err := s.next(); err != nil {
168168
return nil, nil, immediateRetries, errors.Wrapf(err, "failed to get tokens for [%s:%s]", owner.ID(), tokenType)
169169
} else if t == nil {
170170
if !tokensLockedByOthersExist {
@@ -189,9 +189,13 @@ func (s *Selector) selectInternal(ctx context.Context, owner token.OwnerFilter,
189189
}
190190

191191
s.logger.DebugfContext(ctx, "Fetch all non-deleted tokens from the DB and refresh the token cache.")
192-
if s.cache, err = s.fetcher.UnspentTokensIteratorBy(ctx, owner.ID(), tokenType); err != nil {
192+
it, err := s.fetcher.UnspentTokensIteratorBy(ctx, owner.ID(), tokenType)
193+
if err != nil {
193194
return nil, nil, immediateRetries, errors.Wrapf(err, "failed to reload tokens for retry %d [%s:%s]", immediateRetries, owner.ID(), tokenType)
194195
}
196+
if err := s.swapCache(it); err != nil {
197+
return nil, nil, immediateRetries, err
198+
}
195199

196200
immediateRetries++
197201
tokensLockedByOthersExist = false
@@ -222,6 +226,39 @@ func (s *Selector) selectInternal(ctx context.Context, owner token.OwnerFilter,
222226
}
223227
}
224228

229+
// next returns the next token of the current cache. It holds s.mu for the whole
230+
// call so that a concurrent Close cannot swap the iterator out, or nil it, while
231+
// it is being read. It reports an error if the selector has already been closed.
232+
func (s *Selector) next() (*token2.UnspentTokenInWallet, error) {
233+
s.mu.Lock()
234+
defer s.mu.Unlock()
235+
236+
if s.cache == nil {
237+
return nil, errors.New("selector is already closed")
238+
}
239+
240+
return s.cache.Next()
241+
}
242+
243+
// swapCache installs it as the new token cache and closes the iterator it
244+
// replaces, so a refresh on retry does not abandon a database cursor and its
245+
// pooled connection. If the selector was closed in the meantime, it closes it
246+
// too and reports an error: no iterator is ever left unclosed.
247+
func (s *Selector) swapCache(it Iterator[*token2.UnspentTokenInWallet]) error {
248+
s.mu.Lock()
249+
defer s.mu.Unlock()
250+
251+
if s.cache == nil {
252+
it.Close()
253+
254+
return errors.New("selector is already closed")
255+
}
256+
s.cache.Close()
257+
s.cache = it
258+
259+
return nil
260+
}
261+
225262
func (s *Selector) Close() error {
226263
s.mu.Lock()
227264
defer s.mu.Unlock()
Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/*
2+
Copyright IBM Corp. All Rights Reserved.
3+
4+
SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
package sherdlock_test
8+
9+
import (
10+
"context"
11+
"runtime"
12+
"sync"
13+
"sync/atomic"
14+
"testing"
15+
16+
"github.com/LFDT-Panurus/panurus/token/services/selector/sherdlock"
17+
"github.com/LFDT-Panurus/panurus/token/services/selector/sherdlock/mocks"
18+
token2 "github.com/LFDT-Panurus/panurus/token/token"
19+
"github.com/stretchr/testify/assert"
20+
"github.com/stretchr/testify/require"
21+
)
22+
23+
// TestSelectorClosesDisplacedIterators verifies that every iterator the retry
24+
// loop replaces is closed. Previously `selectInternal` overwrote `s.cache` on
25+
// each immediate retry without closing the iterator it displaced, leaking a
26+
// database cursor and its pooled connection per retry. See #2019.
27+
func TestSelectorClosesDisplacedIterators(t *testing.T) {
28+
_, metrics := setupMetricsMocks()
29+
30+
mockFetcher := &mocks.FakeTokenFetcher{}
31+
mockLocker := &mocks.FakeTokenLocker{}
32+
33+
var mu sync.Mutex
34+
var iterators []*mocks.FakeIterator[*token2.UnspentTokenInWallet]
35+
mockFetcher.UnspentTokensIteratorByStub = func(_ context.Context, _ string, _ token2.Type) (sherdlock.Iterator[*token2.UnspentTokenInWallet], error) {
36+
it := &mocks.FakeIterator[*token2.UnspentTokenInWallet]{}
37+
// One token, always locked by someone else, then exhausted: this drives a
38+
// refresh of the cache on every pass through the loop.
39+
it.NextReturnsOnCall(0, &token2.UnspentTokenInWallet{
40+
Id: token2.ID{TxId: "tx1", Index: 0},
41+
Type: "ABC",
42+
Quantity: "100",
43+
}, nil)
44+
it.NextReturnsOnCall(1, nil, nil)
45+
46+
mu.Lock()
47+
defer mu.Unlock()
48+
iterators = append(iterators, it)
49+
50+
return it, nil
51+
}
52+
// Every token appears locked by another process, so the selector exhausts its
53+
// immediate retries and gives up with SufficientButLockedFunds.
54+
mockLocker.TryLockReturns(false, nil)
55+
56+
s := sherdlock.NewSelector(sherdlock.Logger(), mockFetcher, mockLocker, 64, metrics)
57+
_, _, err := s.Select(t.Context(), &unitTestMockOwnerFilter{id: "alice"}, "50", "ABC")
58+
require.Error(t, err)
59+
60+
mu.Lock()
61+
defer mu.Unlock()
62+
require.NotEmpty(t, iterators, "the retry loop must have refreshed the cache at least once")
63+
// All but the last iterator were displaced by a refresh and must be closed.
64+
for i, it := range iterators[:len(iterators)-1] {
65+
assert.Equal(t, 1, it.CloseCallCount(), "displaced iterator %d was leaked", i)
66+
}
67+
// The last one is still the current cache; Close must release it.
68+
require.NoError(t, s.Close())
69+
assert.Equal(t, 1, iterators[len(iterators)-1].CloseCallCount(), "the current iterator must be closed by Close")
70+
}
71+
72+
// TestSelectorCloseDuringRetryIsRaceFree runs Close concurrently with a retrying
73+
// Select. Under `-race` this covers the unsynchronized reads and writes of
74+
// `s.cache` inside `selectInternal`, which could also nil-dereference when Close
75+
// won the race between the closed check and the iterator read. See #2019.
76+
func TestSelectorCloseDuringRetryIsRaceFree(t *testing.T) {
77+
_, metrics := setupMetricsMocks()
78+
79+
mockFetcher := &mocks.FakeTokenFetcher{}
80+
mockLocker := &mocks.FakeTokenLocker{}
81+
82+
var created atomic.Int64
83+
mockFetcher.UnspentTokensIteratorByStub = func(_ context.Context, _ string, _ token2.Type) (sherdlock.Iterator[*token2.UnspentTokenInWallet], error) {
84+
created.Add(1)
85+
it := &mocks.FakeIterator[*token2.UnspentTokenInWallet]{}
86+
it.NextReturnsOnCall(0, &token2.UnspentTokenInWallet{
87+
Id: token2.ID{TxId: "tx1", Index: 0},
88+
Type: "ABC",
89+
Quantity: "100",
90+
}, nil)
91+
it.NextReturnsOnCall(1, nil, nil)
92+
93+
return it, nil
94+
}
95+
mockLocker.TryLockReturns(false, nil)
96+
97+
s := sherdlock.NewSelector(sherdlock.Logger(), mockFetcher, mockLocker, 64, metrics)
98+
99+
var wg sync.WaitGroup
100+
wg.Add(2)
101+
go func() {
102+
defer wg.Done()
103+
// Either outcome is valid: locked funds, or "already closed" if Close won.
104+
_, _, _ = s.Select(t.Context(), &unitTestMockOwnerFilter{id: "alice"}, "50", "ABC")
105+
}()
106+
go func() {
107+
defer wg.Done()
108+
// Give the selector a chance to enter the retry loop before closing it.
109+
for created.Load() == 0 {
110+
runtime.Gosched()
111+
}
112+
_ = s.Close()
113+
}()
114+
wg.Wait()
115+
116+
// Whoever ran last, the selector ends up closed and stays closed.
117+
_, _, err := s.Select(t.Context(), &unitTestMockOwnerFilter{id: "alice"}, "50", "ABC")
118+
require.ErrorContains(t, err, "selector is already closed")
119+
}

0 commit comments

Comments
 (0)