Skip to content

Commit bb6ddf4

Browse files
committed
feat: implement chunked file transfer protocol
This change replaces the simple full-file transfer mechanism with a robust chunked protocol that supports: - Concurrent uploads/downloads with unique transfer IDs - 1MB chunk streaming to support large files - Transfer state management with acknowledgments - Round-robin scheduling for fair resource usage during concurrent transfers Signed-off-by: llogen <christoph.lange@blindspot.software>
1 parent f2b5ea8 commit bb6ddf4

18 files changed

Lines changed: 2041 additions & 342 deletions

File tree

cmds/dutagent/states.go

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -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
154157
func 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

Comments
 (0)