resolved comments

sdks/go/pkg/beam/core/runtime/harness/harness.go

@@ -117,9 +117,9 @@ func Main(ctx context.Context, loggingEndpoint, controlEndpoint string, options
 		if err != nil {
 			log.Errorf(ctx, "error establishing connection to worker status API: %v", err)
 		} else {
-			statusHandler.wg.Add(1)
-			statusHandler.start(ctx)
-			defer statusHandler.stop(ctx)
+			if err := statusHandler.start(ctx); err == nil {
+				defer statusHandler.stop(ctx)
+			}
 		}
 	}
 

sdks/go/pkg/beam/core/runtime/harness/worker_status.go

@@ -52,14 +52,16 @@ func (w *workerStatusHandler) shutdown() {
 }
 
 // start starts the reader to accept WorkerStatusRequest and send WorkerStatusResponse with WorkerStatus API.
-func (w *workerStatusHandler) start(ctx context.Context) {
+func (w *workerStatusHandler) start(ctx context.Context) error {
 	statusClient := fnpb.NewBeamFnWorkerStatusClient(w.conn)
 	stub, err := statusClient.WorkerStatus(ctx)
 	if err != nil {
 		log.Errorf(ctx, "status client not established: %v", err)
-		return
+		return errors.WithContext(err, "status endpoint client not established")
 	}
+	w.wg.Add(1)
 	go w.reader(ctx, stub)
+	return nil
 }
 
 // reader reads the WorkerStatusRequest from the stream and sends a processed WorkerStatusResponse to
@@ -69,7 +71,7 @@ func (w *workerStatusHandler) reader(ctx context.Context, stub fnpb.BeamFnWorker
 	buf := make([]byte, 1<<16)
 	for w.isAlive() {
 		req, err := stub.Recv()
-		if err != nil {
+		if err != nil && err != io.EOF {
 			log.Debugf(ctx, "exiting workerStatusHandler.Reader(): %v", err)
 			return
 		}
@@ -83,10 +85,12 @@ func (w *workerStatusHandler) reader(ctx context.Context, stub fnpb.BeamFnWorker
 }
 
 // stop stops the reader and closes worker status endpoint connection with the runner.
-func (w *workerStatusHandler) stop(ctx context.Context) {
+func (w *workerStatusHandler) stop(ctx context.Context) error {
 	w.shutdown()
 	w.wg.Wait()
 	if err := w.conn.Close(); err != nil {
 		log.Errorf(ctx, "error closing status endpoint connection: %v", err)
+		return errors.WithContext(err, "error closing status endpoint connection")
 	}
+	return nil
 }

sdks/go/pkg/beam/core/runtime/harness/worker_status_test.go

@@ -42,17 +42,17 @@ func (w *BeamFnWorkerStatusServicer) WorkerStatus(b fnpb.BeamFnWorkerStatus_Work
 	return nil
 }
 
-const buffsize = 1024 * 1024
-
 var lis *bufconn.Listener
 
 func setup(t *testing.T, srv *BeamFnWorkerStatusServicer) {
+	const buffsize = 1024 * 1024
 	server := grpc.NewServer()
 	lis = bufconn.Listen(buffsize)
 	fnpb.RegisterBeamFnWorkerStatusServer(server, srv)
 	go func() {
 		if err := server.Serve(lis); err != nil {
 			log.Fatalf("failed to serve: %v", err)
+			panic(err)
 		}
 	}()
 	t.Cleanup(func() {
@@ -68,19 +68,24 @@ func TestSendStatusResponse(t *testing.T) {
 	ctx := context.Background()
 	srv := &BeamFnWorkerStatusServicer{response: make(chan string)}
 	setup(t, srv)
-	conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(dialer), grpc.WithInsecure())
+
+	conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithInsecure(), grpc.WithContextDialer(dialer))
 	if err != nil {
 		t.Fatalf("unable to start test server: %v", err)
 	}
+
 	statusHandler := workerStatusHandler{conn: conn}
-	statusHandler.wg.Add(1)
-	statusHandler.start(ctx)
-	t.Cleanup(func() {
-		statusHandler.stop(ctx)
-	})
+	if err := statusHandler.start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
 	response := []string{}
 	response = append(response, <-srv.response)
 	if len(response) == 0 {
-		t.Errorf("error in response: %v", response)
+		t.Errorf("no response received: %v", response)
+	}
+
+	if err := statusHandler.stop(ctx); err != nil {
+		t.Error(err)
 	}
 }