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) + } +}