From 2118801ccac4f2e8bfcb516e5cb80cf904e353a6 Mon Sep 17 00:00:00 2001 From: Fabian Wienand Date: Thu, 9 Apr 2026 11:46:42 +0200 Subject: [PATCH] fix: unblock session calls when broker workers shut down MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All session methods (Print, Printf, Println, RequestFile, SendFile) blocked indefinitely on unbuffered channels if the companion toClientWorker had already exited — for example when fromClientWorker returned an error first and cancelled the shared worker context. Add a done <-chan struct{} field to session, set to workerCtx.Done() in Broker.Start. Every blocking channel send/receive in the session now selects on done and returns immediately (or with an error for file operations) when the broker is shutting down. Signed-off-by: Fabian Wienand --- internal/dutagent/broker.go | 2 + internal/dutagent/broker_test.go | 87 ++++++++++++++++++++++++++++++++ internal/dutagent/session.go | 40 ++++++++++++--- internal/dutagent/worker.go | 8 ++- 4 files changed, 130 insertions(+), 7 deletions(-) diff --git a/internal/dutagent/broker.go b/internal/dutagent/broker.go index 333deb93..58f1cc81 100644 --- a/internal/dutagent/broker.go +++ b/internal/dutagent/broker.go @@ -53,6 +53,8 @@ func (b *Broker) Start(ctx context.Context, s Stream) (module.Session, <-chan er workerCtx, workerCancel := context.WithCancel(ctx) + b.session.done = workerCtx.Done() + b.wg.Add(numWorkers) b.toClient(workerCtx, workerCancel) b.fromClient(workerCtx, workerCancel) diff --git a/internal/dutagent/broker_test.go b/internal/dutagent/broker_test.go index d8db2c67..2d525679 100644 --- a/internal/dutagent/broker_test.go +++ b/internal/dutagent/broker_test.go @@ -8,6 +8,7 @@ import ( "context" "errors" "io" + "strings" "testing" "time" @@ -276,3 +277,89 @@ func TestBroker_DualErrors(t *testing.T) { // WANT both errors; order unspecified. } } + +// TestSession_PrintNotBlockingOnShutdown verifies that Print, Printf, and Println +// return promptly when the broker workers have shut down (done is closed). +func TestSession_PrintNotBlockingOnShutdown(t *testing.T) { + b := &Broker{} + // recvErrs with a real error triggers fromClientWorker to cancel the context, + // which closes done and causes toClientWorker to exit. + stream := &testStream{recvErrs: []error{errors.New("recv died")}} + sess, errCh := b.Start(context.Background(), stream) + + // Wait for workers to shut down. + collectErrors(t, errCh, 300*time.Millisecond) + + // All three Print variants must return without blocking. + done := make(chan struct{}) + go func() { + defer close(done) + sess.Print("a") + sess.Printf("b %s", "c") + sess.Println("d") + }() + + select { + case <-done: + // ok + case <-time.After(200 * time.Millisecond): + t.Fatal("Print/Printf/Println blocked after broker shutdown") + } +} + +// TestSession_RequestFileErrorOnShutdown verifies that RequestFile returns an error +// when done is closed while the module is waiting for a file to be handed over. +func TestSession_RequestFileErrorOnShutdown(t *testing.T) { + b := &Broker{} + // Block Receive so fromClientWorker doesn't exit on its own; we cancel manually. + stream := &testStream{recvBlock: true} + ctx, cancel := context.WithCancel(context.Background()) + sess, errCh := b.Start(ctx, stream) + + done := make(chan error, 1) + go func() { + _, err := sess.(*session).RequestFile("firmware.bin") + done <- err + }() + + // Give the goroutine time to block on fileReqCh/fileCh. + time.Sleep(20 * time.Millisecond) + + // Cancel to shut down workers (closes done channel on session). + cancel() + if stream.unblockCh != nil { + close(stream.unblockCh) + } + + collectErrors(t, errCh, 300*time.Millisecond) + + select { + case err := <-done: + if err == nil { + t.Fatal("expected RequestFile to return an error on shutdown, got nil") + } + case <-time.After(200 * time.Millisecond): + t.Fatal("RequestFile blocked after broker shutdown") + } +} + +// TestSession_SendFileErrorOnShutdown verifies that SendFile returns an error +// when done is closed before the broker's toClientWorker picks up the file channel. +func TestSession_SendFileErrorOnShutdown(t *testing.T) { + b := &Broker{} + // Use a send error so toClientWorker exits as soon as it tries to send, + // which cancels the context and closes done before SendFile can hand off the file. + stream := &testStream{ + sendErr: errors.New("send died"), + recvErrs: []error{errors.New("recv died")}, + } + sess, errCh := b.Start(context.Background(), stream) + + // Wait for workers to exit. + collectErrors(t, errCh, 300*time.Millisecond) + + err := sess.(*session).SendFile("result.bin", strings.NewReader("data")) + if err == nil { + t.Fatal("expected SendFile to return an error on shutdown, got nil") + } +} diff --git a/internal/dutagent/session.go b/internal/dutagent/session.go index d8ea6222..15b2dfea 100644 --- a/internal/dutagent/session.go +++ b/internal/dutagent/session.go @@ -14,6 +14,7 @@ import ( // session implements the module.Session interface. type session struct { + done <-chan struct{} // closed when broker workers shut down; unblocks pending session calls printCh chan string stdinCh chan []byte stdoutCh chan []byte @@ -28,15 +29,24 @@ type session struct { } func (s *session) Print(a ...any) { - s.printCh <- fmt.Sprint(a...) + select { + case s.printCh <- fmt.Sprint(a...): + case <-s.done: + } } func (s *session) Printf(format string, a ...any) { - s.printCh <- fmt.Sprintf(format, a...) + select { + case s.printCh <- fmt.Sprintf(format, a...): + case <-s.done: + } } func (s *session) Println(a ...any) { - s.printCh <- fmt.Sprintln(a...) + select { + case s.printCh <- fmt.Sprintln(a...): + case <-s.done: + } } //nolint:nonamedreturns @@ -72,9 +82,19 @@ func (s *session) RequestFile(name string) (io.Reader, error) { log.Printf("Module issued file request for: %q", name) - s.fileReqCh <- name // Send the file request to the client. + select { + case s.fileReqCh <- name: + case <-s.done: + return nil, fmt.Errorf("session closed before file request %q could be sent", name) + } + + var file chan []byte - file := <-s.fileCh // This will block until the client sends the file. + select { + case file = <-s.fileCh: + case <-s.done: + return nil, fmt.Errorf("session closed while waiting for file %q", name) + } r, err := chanio.NewChanReader(file) if err != nil { @@ -99,7 +119,15 @@ func (s *session) SendFile(name string, r io.Reader) error { s.currentFile = name file := make(chan []byte, 1) - s.fileCh <- file + + select { + case s.fileCh <- file: + case <-s.done: + s.currentFile = "" + + return fmt.Errorf("session closed before file %q could be sent", name) + } + file <- content close(file) // indicate EOF. diff --git a/internal/dutagent/worker.go b/internal/dutagent/worker.go index 62484fe9..bb97a444 100644 --- a/internal/dutagent/worker.go +++ b/internal/dutagent/worker.go @@ -216,7 +216,13 @@ func fromClientWorker(ctx context.Context, stream Stream, s *session) error { log.Printf("Server received file %q from client", path) file := make(chan []byte, 1) - s.fileCh <- file + + select { + case <-ctx.Done(): + return nil + case s.fileCh <- file: + } + file <- content close(file)