@@ -58,10 +58,11 @@ type Harvester interface {
5858 Run (inputv2.Context , Source , Cursor , Publisher , * Metrics ) error
5959}
6060
61- // reader is the handle for one running harvester's registration in a readerGroup.
61+ // reader is the handle for one harvester's registration in a readerGroup.
6262type reader struct {
63- group * readerGroup
64- srcID string // current registration key, guarded by group.mu
63+ group * readerGroup
64+ // srcID is the current registration key, guarded by group.mu.
65+ srcID string
6566 cancel context.CancelFunc
6667}
6768
@@ -72,21 +73,53 @@ func (rd *reader) currentID() string {
7273 return rd .srcID
7374}
7475
75- // remove cancels the reader's context and deletes its registration from the group.
76+ // currentResource atomically returns the reader's registration key and the store resource for it.
77+ func (rd * reader ) currentResource (store * store ) (string , * resource ) {
78+ rd .group .mu .Lock ()
79+ defer rd .group .mu .Unlock ()
80+
81+ return rd .srcID , store .Get (rd .srcID )
82+ }
83+
84+ // register upgrades a reservation into a running registration and returns the harvester's context.
85+ // It registers under the reader's current ID, which may differ from the reserved one.
86+ func (rd * reader ) register (cancelation inputv2.Canceler ) (context.Context , error ) {
87+ rd .group .mu .Lock ()
88+ defer rd .group .mu .Unlock ()
89+
90+ if rd .group .table [rd .srcID ] != rd {
91+ return nil , fmt .Errorf ("reservation for %s was removed before the harvester started" , rd .srcID )
92+ }
93+ if rd .cancel != nil {
94+ return nil , ErrHarvesterAlreadyRunning
95+ }
96+
97+ ctx , cancel := context .WithCancel (ctxtool .FromCanceller (cancelation ))
98+ rd .cancel = cancel
99+ return ctx , nil
100+ }
101+
102+ // remove cancels the reader's context and deletes its registration from the group. A removed reader
103+ // cannot register anymore.
76104func (rd * reader ) remove () {
77105 rd .group .mu .Lock ()
78106 defer rd .group .mu .Unlock ()
79107
80- rd .cancel ()
108+ rd .unsafeRemove ()
109+ }
110+
111+ func (rd * reader ) unsafeRemove () {
112+ if rd .cancel != nil {
113+ rd .cancel ()
114+ }
81115 if rd .group .table [rd .srcID ] == rd {
82116 delete (rd .group .table , rd .srcID )
83117 }
84118}
85119
86120type readerGroup struct {
87121 mu sync.Mutex
88- // table maps each source ID to its reader handle; a nil value is a
89- // reservation made by reserve before the harvester registers itself.
122+ // table maps each source ID to its non-nil reader handle
90123 table map [string ]* reader
91124}
92125
@@ -96,70 +129,54 @@ func newReaderGroup() *readerGroup {
96129 }
97130}
98131
99- // newContext creates a new context and registers a reader for the given id within the reader group.
100- // The returned reader carries the cancel function for the context; calling it is optional and does
101- // not remove the registration.
102- // An error is returned if a reader is already registered for the id.
103- //
104- // The context will be automatically cancelled once the ID is removed from the group.
105- func (r * readerGroup ) newContext (id string , cancelation inputv2.Canceler ) (context.Context , * reader , error ) {
132+ func (r * readerGroup ) remove (id string ) {
106133 r .mu .Lock ()
107134 defer r .mu .Unlock ()
108135
109- if existing := r .table [id ]; existing != nil {
110- return nil , nil , ErrHarvesterAlreadyRunning
136+ if rd := r .table [id ]; rd != nil {
137+ rd . unsafeRemove ()
111138 }
112-
113- ctx , cancel := context .WithCancel (ctxtool .FromCanceller (cancelation ))
114-
115- rd := & reader {group : r , srcID : id , cancel : cancel }
116- r .table [id ] = rd
117- return ctx , rd , nil
118139}
119140
120- func (r * readerGroup ) remove (id string ) {
141+ // migrate atomically applies updateStore and re-keys any harvester registration from oldID to
142+ // newID, keeping the same reader handle.
143+ func (r * readerGroup ) migrate (oldID , newID string , updateStore func () error ) error {
121144 r .mu .Lock ()
122145 defer r .mu .Unlock ()
123146
124- if rd := r .table [id ]; rd != nil {
125- rd .cancel ()
147+ if _ , exists := r .table [newID ]; exists {
148+ // Target occupied — don't clobber an existing registration.
149+ return fmt .Errorf ("a harvester is already registered for %q" , newID )
126150 }
127151
128- delete (r .table , id )
129- }
130-
131- // migrate moves a running harvester's registration from oldID to newID, keeping the same reader.
132- func (r * readerGroup ) migrate (oldID , newID string ) bool {
133- r .mu .Lock ()
134- defer r .mu .Unlock ()
152+ if err := updateStore (); err != nil {
153+ return err
154+ }
135155
136156 rd := r .table [oldID ]
137157 if rd == nil {
138- // Nothing running under oldID (absent or only reserved).
139- return false
140- }
141-
142- if _ , exists := r .table [newID ]; exists {
143- // Target occupied — don't clobber an existing registration.
144- return false
158+ // No harvester for oldID; only the store needed re-keying.
159+ return nil
145160 }
146161
147162 delete (r .table , oldID )
148163 rd .srcID = newID
149164 r .table [newID ] = rd
150- return true
165+ return nil
151166}
152167
153- func (r * readerGroup ) reserve (id string ) bool {
168+ // reserve creates a reservation for the id, to be upgraded via reader.register once the harvester
169+ // goroutine runs. It returns nil if the id is already taken.
170+ func (r * readerGroup ) reserve (id string ) * reader {
154171 r .mu .Lock ()
155172 defer r .mu .Unlock ()
156173
157- _ , ok := r .table [id ]
158- if ok {
159- return false
174+ if _ , exists := r .table [id ]; exists {
175+ return nil
160176 }
161- r .table [id ] = nil
162- return true
177+ rd := & reader {group : r , srcID : id }
178+ r .table [id ] = rd
179+ return rd
163180}
164181
165182// HarvesterGroup is responsible for running the
@@ -178,7 +195,7 @@ type HarvesterGroup interface {
178195 // SetObserver sets the observer to get notified when a harvester closes
179196 SetObserver (c chan HarvesterStatus )
180197 // Migrate moves a running harvester's bookkeeping registration in-place
181- Migrate (oldID string , next Source )
198+ Migrate (oldID string , next Source , updateStore func ( newID string ) error ) error
182199}
183200
184201type defaultHarvesterGroup struct {
@@ -266,29 +283,28 @@ func startHarvester(
266283 inputID string ,
267284) func (context.Context ) error {
268285 srcID := hg .identifier .ID (src )
269- if ! restart && ! hg .readers .reserve (srcID ) {
270- // A harvester is already running for this source, no need to start another.
271- // This check must happen here, before task.Group.Go spawns a goroutine.
272- // When harvester_limit is set, the spawned goroutine blocks on a semaphore
273- // until a slot is available. Without this early check, repeated file events
274- // would spawn goroutines that wait on the semaphore only to discover (after
275- // acquiring it) that a harvester is already running, causing a goroutine leak.
276- ctx .Logger .Debugf ("Harvester already running for %s" , srcID )
277- return nil
286+ // rd is this harvester's registration handle; all key-dependent steps go through it rather than
287+ // srcID because a migration can re-key the registration at any time.
288+ var rd * reader
289+ if ! restart {
290+ rd = hg .readers .reserve (srcID )
291+ if rd == nil {
292+ // A harvester is already running for this source, no need to start another. This check
293+ // must happen here, before task.Group.Go spawns a goroutine. When harvester_limit is
294+ // set, the spawned goroutine blocks on a semaphore until a slot is available. Without
295+ // this early check, repeated file events would spawn goroutines that wait on the
296+ // semaphore only to discover (after acquiring it) that a harvester is already running.
297+ ctx .Logger .Debugf ("Harvester already running for %s" , srcID )
298+ return nil
299+ }
278300 }
279301
280302 return func (canceler context.Context ) (err error ) {
281- // rd is this harvester's registration handle; cleanup goes through it
282- // rather than srcID because a migration can re-key the registration.
283- var rd * reader
284303 defer func () {
285304 if v := recover (); v != nil {
286305 err = fmt .Errorf ("harvester panic for source %q: %+v\n %s" , srcID , v , debug .Stack ())
287306 if rd != nil {
288307 rd .remove ()
289- } else {
290- // Not registered yet, only the reserve() placeholder exists.
291- hg .readers .remove (srcID )
292308 }
293309 }
294310
@@ -305,27 +321,22 @@ func startHarvester(
305321 ctx .Logger = ctx .Logger .With ("source_file" , srcID )
306322
307323 if restart {
308- // stop previous harvester
324+ // Stop the previous harvester and take its place.
309325 hg .readers .remove (srcID )
326+ rd = hg .readers .reserve (srcID )
327+ if rd == nil {
328+ // Another Start raced in; leave the source to it.
329+ ctx .Logger .Debugf ("Harvester already running for %s" , srcID )
330+ return nil
331+ }
310332 }
311333
312- var harvesterCtx context.Context
313- harvesterCtx , rd , err = hg .readers .newContext (srcID , canceler )
334+ harvesterCtx , err := rd .register (canceler )
314335 if err != nil {
315- // The only possible returned error is ErrHarvesterAlreadyRunning, which is a normal
316- // behaviour of the Filestream input, it's not really an error, it's just a situation.
317- // If the harvester is already running we don't need to start a new one.
318- // At the moment of writing even the returned error is ignored. So the
319- // only real effect of this branch is to not start a second harvester.
320- //
321- // Currently the only places this error is checked is on task.Group and the
322- // only thing it does is to log the error. So to avoid unnecessary errors,
323- // we just return nil.
324- if errors .Is (err , ErrHarvesterAlreadyRunning ) {
325- ctx .Logger .Debug ("Harvester already running" )
326- return nil
327- }
328- return fmt .Errorf ("error while adding new reader to the bookkeeper %w" , err )
336+ // A normal situation, not really an error: the source was stopped before this goroutine
337+ // got to run.
338+ ctx .Logger .Debugf ("Harvester not started: %v" , err )
339+ return nil
329340 }
330341
331342 defer func () {
@@ -339,8 +350,8 @@ func startHarvester(
339350 ctx .Cancelation = harvesterCtx
340351 defer rd .cancel ()
341352
342- resource , err := lock ( ctx , hg .store , srcID )
343- if err != nil {
353+ id , resource := rd . currentResource ( hg .store )
354+ if err := lockResource ( ctx , resource , id ); err != nil {
344355 rd .remove ()
345356 return fmt .Errorf ("error while locking resource: %w" , err )
346357 }
@@ -450,21 +461,17 @@ func (hg *defaultHarvesterGroup) Continue(ctx inputv2.Context, previous, next So
450461 }
451462}
452463
453- // Stop stops the running Harvester for a given Source.
464+ // Stop stops the running (or pending) Harvester for a given Source. It is synchronous so that a
465+ // Stop can never outrun a later Start.
454466func (hg * defaultHarvesterGroup ) Stop (s Source ) {
455- _ = hg .tg .Go (func (_ context.Context ) error {
456- hg .readers .remove (hg .identifier .ID (s ))
457- return nil
458- })
467+ hg .readers .remove (hg .identifier .ID (s ))
459468}
460469
461- // Migrate moves a running harvester's bookkeeping registration from oldID to next's identity
462- // WITHOUT stopping it. It is used after an in-place registry key migration so that a subsequent
463- // Start(next) finds the harvester already registered and no-ops, instead of spawning a duplicate
464- // that would block forever on the shared resource lock. It does nothing if no harvester runs under
465- // oldID or next's key is already taken.
466- func (hg * defaultHarvesterGroup ) Migrate (oldID string , next Source ) {
467- hg .readers .migrate (oldID , hg .identifier .ID (next ))
470+ // Migrate re-keys the registry entry and any harvester registration from oldID to next's identity
471+ // without stopping the harvester.
472+ func (hg * defaultHarvesterGroup ) Migrate (oldID string , next Source , updateStore func (newID string ) error ) error {
473+ newID := hg .identifier .ID (next )
474+ return hg .readers .migrate (oldID , newID , func () error { return updateStore (newID ) })
468475}
469476
470477// StopHarvesters stops all running Harvesters.
@@ -476,23 +483,30 @@ func (hg *defaultHarvesterGroup) StopHarvesters() error {
476483// the cursor state and unlock the key.
477484func lock (ctx inputv2.Context , store * store , key string ) (* resource , error ) {
478485 resource := store .Get (key )
486+ if err := lockResource (ctx , resource , key ); err != nil {
487+ return nil , err
488+ }
489+ return resource , nil
490+ }
479491
492+ // lockResource locks an already retained resource for exclusive access.
493+ func lockResource (ctx inputv2.Context , resource * resource , key string ) error {
480494 if ! resource .lock .TryLock () {
481495 ctx .Logger .Infof ("Resource '%s' currently in use, waiting..." , key )
482496 err := resource .lock .LockContext (ctx .Cancelation )
483497 ctx .Logger .Infof ("Resource '%s' finally released. Lock acquired" , key )
484498 if err != nil {
485499 ctx .Logger .Infof ("Input for resource '%s' has been stopped while waiting" , key )
486500 resource .Release ()
487- return nil , err
501+ return err
488502 }
489503 }
490504
491505 resource .stateMutex .Lock ()
492506 resource .lockedVersion = resource .version
493507 resource .stateMutex .Unlock ()
494508
495- return resource , nil
509+ return nil
496510}
497511
498512func releaseResource (resource * resource ) {
0 commit comments