sdk/dataloader: fix lock around reset()/create new batch logic.

This commit is contained in:
fiatjaf
2025-03-25 19:45:15 -03:00
parent 15d946dad8
commit d1fca24cc3

View File

@@ -87,6 +87,9 @@ func (l *Loader[K, V]) Load(ctx context.Context, key K) (value V, err error) {
// start a sleeper for the current batcher // start a sleeper for the current batcher
l.thresholdReached = make(chan bool) l.thresholdReached = make(chan bool)
// unlock either here or on the else condition
l.batchLock.Unlock()
// we will run the batch function either after some time or after a threshold has been reached // we will run the batch function either after some time or after a threshold has been reached
b := l.curBatcher b := l.curBatcher
go func() { go func() {
@@ -122,6 +125,8 @@ func (l *Loader[K, V]) Load(ctx context.Context, key K) (value V, err error) {
close(req.channel) close(req.channel)
} }
}() }()
} else {
l.batchLock.Unlock()
} }
l.curBatcher.requests = append(l.curBatcher.requests, req) l.curBatcher.requests = append(l.curBatcher.requests, req)
@@ -135,8 +140,6 @@ func (l *Loader[K, V]) Load(ctx context.Context, key K) (value V, err error) {
l.reset() l.reset()
} }
l.batchLock.Unlock()
if v, ok := <-c; ok { if v, ok := <-c; ok {
return v.Data, v.Error return v.Data, v.Error
} }
@@ -145,6 +148,8 @@ func (l *Loader[K, V]) Load(ctx context.Context, key K) (value V, err error) {
} }
func (l *Loader[K, V]) reset() { func (l *Loader[K, V]) reset() {
l.batchLock.Lock()
defer l.batchLock.Unlock()
l.count = 0 l.count = 0
l.curBatcher = nil l.curBatcher = nil
} }