From 5e3b97b44b28fe204d407e853dea8e6fb7bc77f2 Mon Sep 17 00:00:00 2001 From: Copybara Service Date: Wed, 14 Oct 2026 10:03:59 -0700 Subject: [PATCH] Fix client shutdown deadlock when a message buffer is full `Client.Stop` hangs if one of the per-priority buffers for messages to the server is full, i.e. holds 100 messages or 20 MiB that the server hasn't acked yet. This is likely on machines that can't reach the server for a while, e.g. offline laptops, because local services keep sending messages. Then the process only exits when a watchdog terminates it ungracefully: after 10s on SIGTERM, and after a minute on a config reload (SIGHUP). After a graceful `Die` request, it hangs indefinitely. `Stop` waits for the `RetryLoop`s to terminate before it starts acking the messages left in the outbox, but a `RetryLoop` at its limit doesn't read its input until some of its pending messages are acked, so it never notices that the input was closed. `Stop` now starts draining the outbox before it closes the inputs of the `RetryLoop`s and keeps draining until `SortLoop` has returned. The new tests fill a buffer to its count or size limit, without a communicator and with one that can't deliver messages, and check that `Stop` returns. Without the fix, it doesn't return within the tests' 10s timeout. PiperOrigin-RevId: 999443649 --- fleetspeak/src/client/client.go | 32 +++--- fleetspeak/src/client/client_test.go | 165 +++++++++++++++++++++++++++ 2 files changed, 183 insertions(+), 14 deletions(-) diff --git a/fleetspeak/src/client/client.go b/fleetspeak/src/client/client.go index 0c17a1d6..bf4fef57 100644 --- a/fleetspeak/src/client/client.go +++ b/fleetspeak/src/client/client.go @@ -276,26 +276,16 @@ func (c *Client) Stop() { // // - At this point, the communicator is off, so nothing else should be // draining outbox. We do this ourselves and Ack everything so that the - // RetryLoops are guaranteed to terminate. + // RetryLoops are guaranteed to terminate. We have to start draining + // before waiting for the RetryLoops: a RetryLoop with maxBufferCount or + // maxBufferBytes worth of pending messages doesn't read its input, so it + // only notices that the input was closed once some of them are acked. // // - The fake Acks in 1) are safe because the config manager is stopped. // This means that client services are shut down and the Acks will not be // reported outside of this process. - close(c.outLow) - close(c.outMedium) - close(c.outHigh) - c.retryLoopsDone.Wait() - log.Info("Retry loops have terminated.") - - // - Now, no more messages enter outUnsorted. - // - // - We close outUnsorted and drain outbox, to make sure no messages are lost. - // Once these two things are complete, SortLoop will return, and the client - // can be shut down. - done := make(chan struct{}) - close(c.outUnsorted) go func() { for { select { @@ -306,6 +296,20 @@ func (c *Client) Stop() { } } }() + + close(c.outLow) + close(c.outMedium) + close(c.outHigh) + c.retryLoopsDone.Wait() + log.Info("Retry loops have terminated.") + + // - Now, no more messages enter outUnsorted. + // + // - We close outUnsorted and keep draining outbox, to make sure no messages + // are lost. Once these two things are complete, SortLoop will return, and + // the client can be shut down. + + close(c.outUnsorted) c.sortLoopDone.Wait() done <- struct{}{} log.Info("Messages have been drained.") diff --git a/fleetspeak/src/client/client_test.go b/fleetspeak/src/client/client_test.go index cd1140e2..1ee36f69 100644 --- a/fleetspeak/src/client/client_test.go +++ b/fleetspeak/src/client/client_test.go @@ -17,7 +17,9 @@ package client import ( "bytes" "context" + "errors" "fmt" + "io" "os" "os/exec" "path/filepath" @@ -31,6 +33,7 @@ import ( "time" "github.com/google/fleetspeak/fleetspeak/src/client/clienttestutils" + "github.com/google/fleetspeak/fleetspeak/src/client/comms" "github.com/google/fleetspeak/fleetspeak/src/client/config" "github.com/google/fleetspeak/fleetspeak/src/client/service" "github.com/google/fleetspeak/fleetspeak/src/client/stats" @@ -40,6 +43,7 @@ import ( "google.golang.org/protobuf/proto" fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak" + anypb "google.golang.org/protobuf/types/known/anypb" ) var ( @@ -677,3 +681,164 @@ func signServiceConfig(t *testing.T, cfg *fspb.ClientServiceConfig) *fspb.Signed } return &fspb.SignedClientServiceConfig{ServiceConfig: b} } + +// stopTimeout bounds how long the tests below wait for Client.Stop. Stop only +// has to stop in-process components and ack in-memory messages, so it returns +// well within this time unless it deadlocks. +const stopTimeout = 10 * time.Second + +// queueServerMessages passes n messages for the server with the given priority +// and payload size to cl. It returns a counter of the calls to the messages' +// Ack callbacks. +func queueServerMessages(t *testing.T, cl *Client, n int, priority fspb.Message_Priority, payloadSize int) *atomic.Int32 { + t.Helper() + acks := &atomic.Int32{} + for i := range n { + ctx, cancel := context.WithTimeout(t.Context(), stopTimeout) + err := cl.ProcessMessage(ctx, service.AckMessage{ + M: &fspb.Message{ + Destination: &fspb.Address{ServiceName: "RemoteService"}, + Priority: priority, + Data: &anypb.Any{Value: make([]byte, payloadSize)}, + }, + Ack: func() { acks.Add(1) }, + }) + cancel() + if err != nil { + t.Fatalf("ProcessMessage() of message %d/%d failed: %v", i+1, n, err) + } + } + return acks +} + +// stopWithTimeout calls cl.Stop and reports whether it returned within +// stopTimeout. If it did not, Stop keeps running in the background. +func stopWithTimeout(cl *Client) bool { + stopped := make(chan struct{}) + go func() { + cl.Stop() + close(stopped) + }() + select { + case <-stopped: + return true + case <-time.After(stopTimeout): + return false + } +} + +func TestStopAcksPendingMessages(t *testing.T) { + // Without a communicator, nothing acks the messages for the server, so they + // stay pending in the RetryLoop of their priority until Stop acks them. A + // RetryLoop stops reading new messages while it holds maxBufferCount + // messages or maxBufferBytes of data. HIGH isn't covered because the system + // service sends HIGH messages concurrently, which makes the RetryLoop's + // state unpredictable; all three RetryLoops run the same code. + for _, tc := range []struct { + name string + priority fspb.Message_Priority + messages int + payloadSize int + }{ + { + name: "LowBelowLimits", + priority: fspb.Message_LOW, + messages: maxBufferCount - 1, + }, + { + name: "LowAtCountLimit", + priority: fspb.Message_LOW, + messages: maxBufferCount, + }, + { + name: "MediumAtCountLimit", + priority: fspb.Message_MEDIUM, + messages: maxBufferCount, + }, + { + name: "LowAtSizeLimit", + priority: fspb.Message_LOW, + messages: 1, + payloadSize: maxBufferBytes, + }, + } { + t.Run(tc.name, func(t *testing.T) { + cl, err := New(config.Configuration{}, Components{}) + if err != nil { + t.Fatalf("Unable to create client: %v", err) + } + acks := queueServerMessages(t, cl, tc.messages, tc.priority, tc.payloadSize) + + if !stopWithTimeout(cl) { + t.Fatalf("Stop() with %d pending %v messages of %d bytes did not return within %v", tc.messages, tc.priority, tc.payloadSize, stopTimeout) + } + if got := acks.Load(); got != int32(tc.messages) { + t.Errorf("Stop() acked %d of %d pending messages, want all of them acked once", got, tc.messages) + } + }) + } +} + +// offlineCommunicator is a comms.Communicator which takes messages from the +// outbox but never delivers them, like a communicator which can't reach the +// server. Like the HTTPS communicators, it nacks the messages it holds when it +// is stopped. +type offlineCommunicator struct { + comms.Communicator + + cctx comms.Context + stop chan struct{} + done chan struct{} +} + +func (c *offlineCommunicator) Setup(cctx comms.Context) error { + c.cctx = cctx + return nil +} + +func (c *offlineCommunicator) Start() error { + c.stop = make(chan struct{}) + c.done = make(chan struct{}) + go func() { + defer close(c.done) + var held []comms.MessageInfo + for { + select { + case m := <-c.cctx.Outbox(): + held = append(held, m) + case <-c.stop: + for _, m := range held { + m.Nack() + } + return + } + } + }() + return nil +} + +func (c *offlineCommunicator) Stop() { + close(c.stop) + <-c.done +} + +func (c *offlineCommunicator) GetFileIfModified(context.Context, string, string, time.Time) (io.ReadCloser, time.Time, error) { + return nil, time.Time{}, errors.New("server unreachable") +} + +func TestStopWhileServerUnreachable(t *testing.T) { + cl, err := New(config.Configuration{}, Components{Communicator: &offlineCommunicator{}}) + if err != nil { + t.Fatalf("Unable to create client: %v", err) + } + // Fill the LOW RetryLoop with messages which the communicator takes but + // can't deliver, as happens when a machine is offline for a while. + queueServerMessages(t, cl, maxBufferCount, fspb.Message_LOW, 0) + + // Unlike TestStopAcksPendingMessages, don't count the acks: Stop only acks + // the messages which the RetryLoop resends after the communicator nacks + // them, and the RetryLoop may terminate before resending all of them. + if !stopWithTimeout(cl) { + t.Fatalf("Stop() with %d undeliverable messages did not return within %v", maxBufferCount, stopTimeout) + } +}