Add connections and buffers flush support in polling and streaming clients.

PiperOrigin-RevId: 913782444

Author: Copybara Service

fleetspeak/src/client/client.go

@@ -310,3 +310,23 @@ func (c *Client) Stop() {
 	done <- struct{}{}
 	log.Info("Messages have been drained.")
 }
+
+// FlushServices triggers all services that support it to flush their data.
+func (c *Client) FlushServices(ctx context.Context) {
+	log.InfoContextf(ctx, "Client: FlushServices called")
+	c.sc.FlushServices(ctx)
+}
+
+// ForceResetAndFlushCommunicator programmatically and forcefully tears down the existing
+// network connection, establishes a fresh one, and blocks until the outbox is
+// successfully drained or the context expires.
+func (c *Client) ForceResetAndFlushCommunicator(ctx context.Context) error {
+	log.InfoContextf(ctx, "ForceResetAndFlushCommunicator triggered")
+	if c.com == nil {
+		return fmt.Errorf("communicator not set")
+	}
+	log.InfoContextf(ctx, "Calling Reset")
+	c.com.Reset()
+	log.InfoContextf(ctx, "Reset finished, calling Flush")
+	return c.com.Flush(ctx)
+}

fleetspeak/src/client/comms/comms.go

@@ -39,6 +39,16 @@ type Communicator interface {
 	Start() error        // Tells the communicator to start sending and receiving messages.
 	Stop()               // Tells the communicator to stop sending and receiving messages.
 
+	// Reset forcefully tears down any underlying network connections and
+	// immediately initiates establishing a fresh connection (or polling,
+	// depending on the implementation) to resume communication.
+	Reset()
+
+	// Flush blocks until the communicator has successfully established a
+	// fresh connection and flushed all pending messages, or until the provided
+	// context expires.
+	Flush(ctx context.Context) error
+
 	// GetFileIfModified attempts to retrieve a file from a server, if it
 	// has been modified since modSince. If it has not been modified, it
 	// returns nil. Otherwise, it returns a ReadCloser for the file's data

fleetspeak/src/client/https/polling.go

@@ -58,6 +58,11 @@ type Communicator struct {
 	clientCertificateHeader string
 
 	certBytes []byte
+
+	wakeUp       chan struct{}
+	pollComplete chan error
+	mu           sync.Mutex
+	pollCancel   context.CancelFunc
 }
 
 // Setup implements comms.Communicator.
@@ -108,6 +113,8 @@ func (c *Communicator) configure() error {
 	}
 	c.ctx, c.done = context.WithCancel(context.Background())
 	c.clientCertificateHeader = si.ClientCertificateHeader
+	c.wakeUp = make(chan struct{}, 1)
+	c.pollComplete = make(chan error, 1)
 	c.certBytes = certBytes
 	return nil
 }
@@ -127,6 +134,39 @@ func (c *Communicator) Stop() {
 	c.wd.Stop()
 }
 
+// Reset implements comms.Communicator.
+func (c *Communicator) Reset() {
+	log.Infof("Reset called")
+	c.mu.Lock()
+	if c.pollCancel != nil {
+		c.pollCancel()
+	}
+	c.mu.Unlock()
+	c.hc.Transport.(*http.Transport).CloseIdleConnections()
+	// Drain pollComplete to ensure Flush waits for a new poll.
+	select {
+	case <-c.pollComplete:
+	default:
+	}
+	select {
+	case c.wakeUp <- struct{}{}:
+	default:
+	}
+}
+
+// Flush implements comms.Communicator.
+func (c *Communicator) Flush(ctx context.Context) error {
+	log.InfoContextf(ctx, "Flush called")
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case err := <-c.pollComplete:
+			return err
+		}
+	}
+}
+
 // processingLoop polls the server according to the configured policies while
 // the communicator is active.
 //
@@ -152,11 +192,19 @@ func (c *Communicator) processingLoop() {
 	// and updates the variables defined above. In case of failure it also sleeps
 	// for the MinFailureDelay.
 	poll := func() {
+		var err error
+		defer func() {
+			select {
+			case c.pollComplete <- err:
+			default:
+			}
+		}()
 		c.wd.Reset()
 		if c.cctx.CurrentID() != c.id {
 			c.configure()
 		}
-		active, err := c.poll(toSend)
+		var active bool
+		active, err = c.poll(toSend)
 		if err != nil {
 			log.Warningf("Failure during polling: %v", err)
 			for _, m := range toSend {
@@ -175,6 +223,8 @@ func (c *Communicator) processingLoop() {
 			case <-t.C:
 			case <-c.ctx.Done():
 				t.Stop()
+			case <-c.wakeUp:
+				t.Stop()
 			}
 			return
 		}
@@ -255,6 +305,9 @@ func (c *Communicator) processingLoop() {
 			return
 		case <-t.C:
 			poll()
+		case <-c.wakeUp:
+			t.Stop()
+			poll()
 		case m := <-c.cctx.Outbox():
 			t.Stop()
 			toSend = append(toSend, m)
@@ -324,12 +377,16 @@ func (c *Communicator) pollHost(host string, data []byte) (*fspb.ContactData, er
 
 	u := url.URL{Scheme: "https", Host: host, Path: "/message"}
 
+	c.mu.Lock()
+	compression := c.conf.GetCompression()
+	c.mu.Unlock()
+
 	body := &bytes.Buffer{}
-	if c.conf.GetCompression() == fspb.CompressionAlgorithm_COMPRESSION_NONE {
+	if compression == fspb.CompressionAlgorithm_COMPRESSION_NONE {
 		// Shortcut to prevent unnecessary copying of data
 		body = bytes.NewBuffer(data)
 	} else {
-		bw := CompressingWriter(body, c.conf.GetCompression())
+		bw := CompressingWriter(body, compression)
 		bw.Write(data)
 		bw.Close()
 	}
@@ -339,6 +396,17 @@ func (c *Communicator) pollHost(host string, data []byte) (*fspb.ContactData, er
 	if sendErr != nil {
 		return nil, sendErr
 	}
+	var reqCtx context.Context
+	c.mu.Lock()
+	reqCtx, c.pollCancel = context.WithCancel(c.ctx)
+	c.mu.Unlock()
+	defer func() {
+		c.mu.Lock()
+		c.pollCancel()
+		c.pollCancel = nil
+		c.mu.Unlock()
+	}()
+	req = req.WithContext(reqCtx)
 	SetContentEncoding(req.Header, c.conf.GetCompression())
 	if c.clientCertificateHeader != "" {
 		bc := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: c.certBytes})

fleetspeak/src/client/https/polling_test.go

@@ -733,3 +733,146 @@ func TestCertificateRevoked(t *testing.T) {
 	tl.Close()
 	cl.Stop()
 }
+
+func TestReset(t *testing.T) {
+	var c Communicator
+	conf := config.Configuration{
+		Servers: []string{"localhost:1234"},
+		CommunicatorConfig: &clpb.CommunicatorConfig{
+			MaxPollDelaySeconds:    10,
+			MinFailureDelaySeconds: 10,
+		},
+	}
+	dialCalls := make(chan struct{}, 10)
+	c.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
+		dialCalls <- struct{}{}
+		<-ctx.Done()
+		return nil, ctx.Err()
+	}
+
+	cl, err := client.New(
+		conf,
+		client.Components{
+			Communicator: &c,
+		})
+	if err != nil {
+		t.Fatalf("unable to create client: %v", err)
+	}
+	defer cl.Stop()
+
+	// Wait for first dial attempt
+	select {
+	case <-dialCalls:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Timed out waiting for first dial attempt")
+	}
+
+	// Call Reset
+	time.Sleep(100 * time.Millisecond)
+	c.Reset()
+
+	// Wait for second dial attempt (should be immediate, not waiting 10s)
+	select {
+	case <-dialCalls:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Timed out waiting for second dial attempt after Reset")
+	}
+}
+
+func TestFlush(t *testing.T) {
+	// Create a local https server for the client to talk to.
+	pemCert, pemKey, err := common_util.ServerCert()
+	if err != nil {
+		t.Fatal(err)
+	}
+	cp, err := tls.X509KeyPair(pemCert, pemKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ad, err := net.ResolveTCPAddr("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	tl, err := net.ListenTCP("tcp", ad)
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr := tl.Addr().String()
+
+	pollCount := int32(0)
+	mux := http.NewServeMux()
+	mux.HandleFunc("/message", func(res http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&pollCount, 1)
+		cd := fspb.ContactData{
+			SequencingNonce: 42,
+		}
+		buf, err := proto.Marshal(&cd)
+		if err != nil {
+			t.Fatalf("unable to marshal ContactData: %v", err)
+		}
+		res.Header().Set("Content-Type", "application/octet-stream")
+		res.WriteHeader(http.StatusOK)
+		res.Write(buf)
+	})
+
+	server := http.Server{
+		Addr:    addr,
+		Handler: mux,
+		TLSConfig: &tls.Config{
+			ClientAuth:   tls.RequireAnyClientCert,
+			Certificates: []tls.Certificate{cp},
+			NextProtos:   []string{"h2"},
+		},
+	}
+	l := tls.NewListener(tl, server.TLSConfig)
+	go server.Serve(l)
+	defer tl.Close()
+
+	var c Communicator
+	conf := config.Configuration{
+		TrustedCerts: x509.NewCertPool(),
+		Servers:      []string{addr},
+		CommunicatorConfig: &clpb.CommunicatorConfig{
+			MaxPollDelaySeconds:    10,
+			MinFailureDelaySeconds: 10,
+		},
+	}
+	if !conf.TrustedCerts.AppendCertsFromPEM(pemCert) {
+		t.Fatal("unable to add server cert to pool")
+	}
+
+	cl, err := client.New(
+		conf,
+		client.Components{
+			Communicator: &c,
+		})
+	if err != nil {
+		t.Fatalf("unable to create client: %v", err)
+	}
+	defer cl.Stop()
+
+	// Wait for first poll
+	for atomic.LoadInt32(&pollCount) == 0 {
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	// Call Flush in a goroutine
+	flushDone := make(chan error, 1)
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		flushDone <- c.Flush(ctx)
+	}()
+
+	// Trigger a poll via wakeUp (indirectly via Reset)
+	c.Reset()
+
+	select {
+	case err := <-flushDone:
+		if err != nil {
+			t.Errorf("Flush failed: %v", err)
+		}
+	case <-time.After(6 * time.Second):
+		t.Fatal("Timed out waiting for Flush")
+	}
+}