@@ -19,6 +19,7 @@ import (
1919 "encoding/binary"
2020 "encoding/hex"
2121 "io"
22+ "sync"
2223 "time"
2324
2425 "github.com/bsv-blockchain/go-bt/v2/chainhash"
@@ -47,10 +48,12 @@ type Batcher struct {
4748 writeKeys bool
4849 // queue is a lock-free queue for storing batch items to be processed asynchronously
4950 queue * lockfreequeue.LockFreeQ [BatchItem ]
50- // queueCtx is the context for controlling the background batch processing goroutine
51- queueCtx context.Context
52- // queueCancel is the function to cancel the queue context and stop background processing
53- queueCancel context.CancelFunc
51+ // done is the channel for signaling the background batch processing goroutine to stop
52+ done chan struct {}
53+ // notifyCh is used to notify the worker goroutine when new items are enqueued
54+ notifyCh chan struct {}
55+ // wg is used to wait for the background worker goroutine to complete during shutdown
56+ wg sync.WaitGroup
5457 // currentBatch holds the accumulated blob data for the current batch
5558 currentBatch []byte
5659 // currentBatchKeys holds the accumulated key data for the current batch (if writeKeys is true)
@@ -94,28 +97,41 @@ type blobStoreSetter interface {
9497// Returns:
9598// - *Batcher: A configured batcher instance ready to accept blob operations
9699func New (logger ulogger.Logger , blobStore blobStoreSetter , sizeInBytes int , writeKeys bool ) * Batcher {
97- ctx , cancel := context .WithCancel (context .Background ())
98100 b := & Batcher {
99101 logger : logger ,
100102 blobStore : blobStore ,
101103 sizeInBytes : sizeInBytes ,
102104 writeKeys : writeKeys ,
103105 queue : lockfreequeue .NewLockFreeQ [BatchItem ](),
104- queueCtx : ctx ,
105- queueCancel : cancel ,
106+ done : make ( chan struct {}) ,
107+ notifyCh : make ( chan struct {}, 1 ) ,
106108 currentBatch : make ([]byte , 0 , sizeInBytes ),
107109 currentBatchKeys : make ([]byte , 0 , sizeInBytes ),
108110 }
109111
112+ b .wg .Add (1 )
110113 go func () {
114+ defer b .wg .Done ()
115+
111116 var (
112117 batchItem * BatchItem
113118 err error
114119 )
115120
116121 for {
122+ // Try immediate dequeue (optimistic fast path)
123+ batchItem = b .queue .Dequeue ()
124+ if batchItem != nil {
125+ err = b .processBatchItem (batchItem )
126+ if err != nil {
127+ b .logger .Errorf ("error processing batch item: %v" , err )
128+ }
129+ continue
130+ }
131+
132+ // Queue is empty - wait for notification or shutdown
117133 select {
118- case <- b .queueCtx . Done () :
134+ case <- b .done :
119135 // Process remaining items before exiting
120136 for {
121137 batchItem = b .queue .Dequeue ()
@@ -133,19 +149,11 @@ func New(logger ulogger.Logger, blobStore blobStoreSetter, sizeInBytes int, writ
133149 b .logger .Errorf ("error writing final batch during shutdown: %v" , err )
134150 }
135151 }
136-
137152 return
138- default :
139- batchItem = b .queue .Dequeue ()
140- if batchItem == nil {
141- time .Sleep (10 * time .Millisecond )
142- continue
143- }
144153
145- err = b .processBatchItem (batchItem )
146- if err != nil {
147- b .logger .Errorf ("error processing batch item: %v" , err )
148- }
154+ case <- b .notifyCh :
155+ // Item available, loop back to dequeue
156+ continue
149157 }
150158 }
151159 }()
@@ -235,7 +243,9 @@ func (b *Batcher) writeBatch(currentBatch []byte, batchKeys []byte) error {
235243 binary .BigEndian .PutUint32 (batchKey , timeUint32 )
236244 // add a random string as the next bytes, to prevent conflicting filenames from other pods
237245 randBytes := make ([]byte , 4 )
238- _ , _ = rand .Read (randBytes )
246+ if _ , err := rand .Read (randBytes ); err != nil {
247+ return errors .NewStorageError ("failed to generate random bytes for batch key" , err )
248+ }
239249 batchKey = append (batchKey , randBytes ... )
240250
241251 g , gCtx := errgroup .WithContext (context .Background ())
@@ -298,10 +308,16 @@ func (b *Batcher) Health(ctx context.Context, checkLiveness bool) (int, string,
298308// - error: Any error that occurred during shutdown
299309func (b * Batcher ) Close (_ context.Context ) error {
300310 // Signal the background goroutine to stop
301- b . queueCancel ( )
311+ close ( b . done )
302312
303- // Wait a bit to ensure the goroutine has time to process remaining items
304- time .Sleep (100 * time .Millisecond )
313+ // Wake up the worker if it's blocked on notifyCh
314+ select {
315+ case b .notifyCh <- struct {}{}:
316+ default :
317+ }
318+
319+ // Wait for the background goroutine to finish processing all remaining items
320+ b .wg .Wait ()
305321
306322 return nil
307323}
@@ -351,6 +367,12 @@ func (b *Batcher) Set(_ context.Context, hash []byte, fileType fileformat.FileTy
351367 value : value ,
352368 })
353369
370+ // Notify worker that new item is available (non-blocking)
371+ select {
372+ case b .notifyCh <- struct {}{}:
373+ default : // Already notified, don't block
374+ }
375+
354376 return nil
355377}
356378
0 commit comments