@@ -31,6 +31,7 @@ type runCmdArgs struct {
3131 cmdMsg * pb.Command
3232 dev dut.Device
3333 cmd dut.Command
34+ broker * dutagent.Broker
3435 session module.Session
3536 moduleErrCh chan error
3637 brokerErrCh <- chan error
@@ -151,8 +152,10 @@ func releaseAutoLock(_ context.Context, args runCmdArgs) (runCmdArgs, fsm.State[
151152// in a separate goroutine, this state will not wait for the execution to finish.
152153// Further, worker goroutines will be started to serve the module-to-client communication
153154// during the module execution.
155+ //
156+ //nolint:funlen
154157func executeModules (ctx context.Context , args runCmdArgs ) (runCmdArgs , fsm.State [runCmdArgs ], error ) {
155- broker : = & dutagent.Broker {}
158+ args . broker = & dutagent.Broker {}
156159
157160 // Deferred initialization of the moduleErr channel: only create if not already provided
158161 // (tests may still pass a custom channel).
@@ -163,7 +166,7 @@ func executeModules(ctx context.Context, args runCmdArgs) (runCmdArgs, fsm.State
163166 rpcCtx := ctx
164167 modCtx , modCtxCancel := context .WithCancel (rpcCtx )
165168
166- moduleSession , brokerErrCh := broker .Start (modCtx , args .stream )
169+ moduleSession , brokerErrCh := args . broker .Start (modCtx , args .stream )
167170 args .brokerErrCh = brokerErrCh
168171 args .session = moduleSession
169172
@@ -178,12 +181,19 @@ func executeModules(ctx context.Context, args runCmdArgs) (runCmdArgs, fsm.State
178181 // Run the modules in a goroutine.
179182 // Termination of the module execution is signaled by closing the moduleErrCh channel.
180183 go func () {
184+ defer modCtxCancel () // Ensure workers exit even if stream doesn't close
185+
181186 cnt := len (args .cmd .Modules )
182187
183188 for idx , module := range args .cmd .Modules {
184189 if ctx .Err () != nil {
185190 log .Printf ("Execution aborted, %d of %d modules done: %v" , idx , cnt , ctx .Err ())
186- modCtxCancel ()
191+ args .broker .Shutdown ()
192+
193+ // Wait for file transfers to complete (workers will exit gracefully)
194+ log .Print ("Waiting for file transfers to complete..." )
195+ args .broker .WaitForTransfersToComplete ()
196+ log .Print ("All file transfers completed" )
187197
188198 return
189199 }
@@ -195,14 +205,25 @@ func executeModules(ctx context.Context, args runCmdArgs) (runCmdArgs, fsm.State
195205 args .moduleErrCh <- err
196206
197207 log .Printf ("Module %q failed: %v" , module .Config .Name , err )
198- modCtxCancel ()
208+ args .broker .Shutdown ()
209+
210+ // Wait for file transfers to complete (workers will exit gracefully)
211+ log .Print ("Waiting for file transfers to complete..." )
212+ args .broker .WaitForTransfersToComplete ()
213+ log .Print ("All file transfers completed" )
199214
200215 return
201216 }
202217 }
203218
204219 log .Print ("All modules finished successfully" )
205- modCtxCancel ()
220+ args .broker .Shutdown ()
221+
222+ // Wait for file transfers to complete (workers will exit gracefully)
223+ log .Print ("Waiting for file transfers to complete..." )
224+ args .broker .WaitForTransfersToComplete ()
225+ log .Print ("All file transfers completed" )
226+
206227 close (args .moduleErrCh )
207228 }()
208229
0 commit comments