[flyteagent] Remove redundant code in Agent (#5454)

* Remove Redundant code in Agent Signed-off-by: Future-Outlier <[email protected]> * improve error messages Signed-off-by: Future-Outlier <[email protected]> * change to Signed-off-by: Future-Outlier <[email protected]> * Remove defer func() by pingsu's advice. Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: pingsutw <[email protected]> --------- Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: pingsutw <[email protected]>
flyteorg · Jun 14, 2024 · bba8c11 · bba8c11
1 parent 9b78125
commit bba8c11
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 17 deletions.
diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/client.go b/flyteplugins/go/tasks/plugins/webapi/agent/client.go
@@ -62,19 +62,12 @@ func getGrpcConnection(ctx context.Context, agent *Deployment) (*grpc.ClientConn
 	if err != nil {
 		return nil, err
 	}
-	defer func() {
-		if err != nil {
-			if cerr := conn.Close(); cerr != nil {
-				grpclog.Infof("Failed to close conn to %s: %v", agent, cerr)
-			}
-			return
+
+	go func() {
+		<-ctx.Done()
+		if cerr := conn.Close(); cerr != nil {
+			grpclog.Infof("Failed to close conn to %s: %v", agent, cerr)
 		}
-		go func() {
-			<-ctx.Done()
-			if cerr := conn.Close(); cerr != nil {
-				grpclog.Infof("Failed to close conn to %s: %v", agent, cerr)
-			}
-		}()
 	}()
 
 	return conn, nil

diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go
@@ -183,6 +183,7 @@ func (p Plugin) ExecuteTaskSync(
 
 	in, err := stream.Recv()
 	if err != nil {
+		logger.Errorf(ctx, "Failed to write output with err %s", err.Error())
 		return nil, nil, err
 	}
 	if in.GetHeader() == nil {
@@ -193,11 +194,7 @@ func (p Plugin) ExecuteTaskSync(
 	resource := in.GetHeader().GetResource()
 
 	if err := stream.CloseSend(); err != nil {
-		return nil, nil, err
-	}
-
-	if err != nil {
-		logger.Errorf(ctx, "Failed to write output with err %s", err.Error())
+		logger.Errorf(ctx, "Failed to close stream with err %s", err.Error())
 		return nil, nil, err
 	}