fix(agent): Graceful Stop (#3219)

* fix(agent): Graceful Stop * fix(agent): Graceful Stop
kubeshop · Oct 5, 2023 · ecbbdcb · ecbbdcb
1 parent e034bce
commit ecbbdcb
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 9 deletions.
diff --git a/agent/client/client.go b/agent/client/client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/kubeshop/tracetest/agent/proto"
@@ -38,6 +39,7 @@ func (c *Client) Start(ctx context.Context) error {
 		return err
 	}
 
+	c.done = make(chan bool)
 	ctx, cancel := context.WithCancel(ctx)
 	go func() {
 		<-c.done
@@ -138,3 +140,7 @@ func (c *Client) getName() (string, error) {
 
 	return hostname, nil
 }
+
+func isCancelledError(err error) bool {
+	return err != nil && strings.Contains(err.Error(), "context canceled")
+}
diff --git a/agent/client/workflow_listen_for_ds_connection_tests.go b/agent/client/workflow_listen_for_ds_connection_tests.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -21,12 +22,12 @@ func (c *Client) startDataStoreConnectionTestListener(ctx context.Context) error
 		for {
 			req := proto.DataStoreConnectionTestRequest{}
 			err := stream.RecvMsg(&req)
-			if err == io.EOF {
+			if errors.Is(err, io.EOF) || isCancelledError(err) {
 				return
 			}
 
 			if err != nil {
-				log.Fatal("could not get message from trigger stream: %w", err)
+				log.Fatal("could not get message from ds connection stream: %w", err)
 			}
 
 			// TODO: Get ctx from request

diff --git a/agent/client/workflow_listen_for_poll_requests.go b/agent/client/workflow_listen_for_poll_requests.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -21,12 +22,12 @@ func (c *Client) startPollerListener(ctx context.Context) error {
 		for {
 			resp := proto.PollingRequest{}
 			err := stream.RecvMsg(&resp)
-			if err == io.EOF {
+			if errors.Is(err, io.EOF) || isCancelledError(err) {
 				return
 			}
 
 			if err != nil {
-				log.Fatal("could not get message from trigger stream: %w", err)
+				log.Fatal("could not get message from polling stream: %w", err)
 			}
 
 			// TODO: Get ctx from request

diff --git a/agent/client/workflow_listen_for_trigger_requests.go b/agent/client/workflow_listen_for_trigger_requests.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -21,7 +22,7 @@ func (c *Client) startTriggerListener(ctx context.Context) error {
 		for {
 			resp := proto.TriggerRequest{}
 			err := stream.RecvMsg(&resp)
-			if err == io.EOF {
+			if errors.Is(err, io.EOF) || isCancelledError(err) {
 				return
 			}
 

diff --git a/agent/client/workflow_shutdown.go b/agent/client/workflow_shutdown.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -21,12 +22,13 @@ func (c *Client) startShutdownListener(ctx context.Context) error {
 		for {
 			resp := proto.ShutdownRequest{}
 			err := stream.RecvMsg(&resp)
-			if err == io.EOF {
+
+			if errors.Is(err, io.EOF) || isCancelledError(err) {
 				return
 			}
 
 			if err != nil {
-				log.Fatal("could not get message from trigger stream: %w", err)
+				log.Fatal("could not get shutdown listener: %w", err)
 			}
 
 			// TODO: get context from request

diff --git a/cli/config/selector.go b/cli/config/selector.go
@@ -34,7 +34,7 @@ func (c Configurator) organizationSelector(ctx context.Context, cfg Config) (str
 	options := make([]cliUI.Option, len(elements))
 	for i, org := range elements {
 		options[i] = cliUI.Option{
-			Text: org.Name,
+			Text: fmt.Sprintf("%s (%s)", org.Name, org.ID),
 			Fn: func(o Entry) func(ui cliUI.UI) {
 				return func(ui cliUI.UI) {
 					orgID = o.ID
@@ -72,7 +72,7 @@ func (c Configurator) environmentSelector(ctx context.Context, cfg Config) (stri
 	options := make([]cliUI.Option, len(elements))
 	for i, env := range elements {
 		options[i] = cliUI.Option{
-			Text: env.Name,
+			Text: fmt.Sprintf("%s (%s)", env.Name, env.ID),
 			Fn: func(e Entry) func(ui cliUI.UI) {
 				return func(ui cliUI.UI) {
 					envID = e.ID