feat: enable agent handling of stop requests (#3539)

agent/client/client.go

@@ -44,6 +44,7 @@ type Client struct {
 
 	logger *zap.Logger
 
+	stopListener                func(context.Context, *proto.StopRequest) error
 	triggerListener             func(context.Context, *proto.TriggerRequest) error
 	pollListener                func(context.Context, *proto.PollingRequest) error
 	shutdownListener            func(context.Context, *proto.ShutdownRequest) error
@@ -66,6 +67,11 @@ func (c *Client) Start(ctx context.Context) error {
 		cancel()
 	}()
 
+	err = c.startStopListener(ctx)
+	if err != nil {
+		return err
+	}
+
 	err = c.startTriggerListener(ctx)
 	if err != nil {
 		return err
@@ -112,6 +118,10 @@ func (c *Client) Close() error {
 	return c.conn.Close()
 }
 
+func (c *Client) OnStopRequest(listener func(context.Context, *proto.StopRequest) error) {
+	c.stopListener = listener
+}
+
 func (c *Client) OnTriggerRequest(listener func(context.Context, *proto.TriggerRequest) error) {
 	c.triggerListener = listener
 }

agent/client/mocks/grpc_server.go

@@ -88,6 +88,14 @@ func (s *GrpcServerMock) Connect(ctx context.Context, req *proto.ConnectRequest)
 	}, nil
 }
 
+func (s *GrpcServerMock) RegisterStopRequestAgent(id *proto.AgentIdentification, stream proto.Orchestrator_RegisterStopRequestAgentServer) error {
+	if id.Token != "token" {
+		return fmt.Errorf("could not validate token")
+	}
+
+	return nil
+}
+
 func (s *GrpcServerMock) RegisterTriggerAgent(id *proto.AgentIdentification, stream proto.Orchestrator_RegisterTriggerAgentServer) error {
 	if id.Token != "token" {
 		return fmt.Errorf("could not validate token")

agent/client/workflow_listen_for_stop_requests.go

@@ -0,0 +1,47 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/kubeshop/tracetest/agent/proto"
+)
+
+func (c *Client) startStopListener(ctx context.Context) error {
+	client := proto.NewOrchestratorClient(c.conn)
+
+	stream, err := client.RegisterStopRequestAgent(ctx, c.sessionConfig.AgentIdentification)
+	if err != nil {
+		return fmt.Errorf("could not open agent stream: %w", err)
+	}
+
+	go func() {
+		for {
+			resp := proto.StopRequest{}
+			err := stream.RecvMsg(&resp)
+			if isEndOfFileError(err) || isCancelledError(err) {
+				return
+			}
+
+			reconnected, err := c.handleDisconnectionError(err)
+			if reconnected {
+				return
+			}
+
+			if err != nil {
+				log.Println("could not get message from trigger stream: %w", err)
+				time.Sleep(1 * time.Second)
+				continue
+			}
+
+			// TODO: get context from request
+			err = c.stopListener(context.Background(), &resp)
+			if err != nil {
+				fmt.Println(err.Error())
+			}
+		}
+	}()
+	return nil
+}

agent/event/observer.go

@@ -18,6 +18,9 @@ type Observer interface {
 	StartDataStoreConnection(*proto.DataStoreConnectionTestRequest)
 	EndDataStoreConnection(*proto.DataStoreConnectionTestRequest, error)
 
+	StartStopRequest(*proto.StopRequest)
+	EndStopRequest(*proto.StopRequest, error)
+
 	Error(error)
 }
 
@@ -110,3 +113,19 @@ func (o *wrapperObserver) EndTriggerExecution(request *proto.TriggerRequest, err
 
 	o.wrappedObserver.EndTriggerExecution(request, err)
 }
+
+func (o *wrapperObserver) StartStopRequest(request *proto.StopRequest) {
+	if o.wrappedObserver == nil {
+		return
+	}
+
+	o.wrappedObserver.StartStopRequest(request)
+}
+
+func (o *wrapperObserver) EndStopRequest(request *proto.StopRequest, err error) {
+	if o.wrappedObserver == nil {
+		return
+	}
+
+	o.wrappedObserver.EndStopRequest(request, err)
+}