Correctly synchronize streams (Fix: some lines cut-off at the end of the output)

Author: cmaglie

arduino/utils/stream.go

@@ -16,15 +16,18 @@
 package utils
 
 import (
+	"context"
 	"io"
 	"time"
 )
 
 // FeedStreamTo creates a pipe to pass data to the writer function.
 // FeedStreamTo returns the io.Writer side of the pipe, on which the user can write data
-func FeedStreamTo(writer func(data []byte)) io.Writer {
+func FeedStreamTo(writer func(data []byte)) (io.WriteCloser, context.Context) {
+	ctx, cancel := context.WithCancel(context.Background())
 	r, w := io.Pipe()
 	go func() {
+		defer cancel()
 		data := make([]byte, 16384)
 		for {
 			if n, err := r.Read(data); err == nil {
@@ -41,7 +44,7 @@ func FeedStreamTo(writer func(data []byte)) io.Writer {
 			}
 		}
 	}()
-	return w
+	return w, ctx
 }
 
 // ConsumeStreamFrom creates a pipe to consume data from the reader function.

commands/daemon/daemon.go

@@ -266,15 +266,19 @@ func (s *ArduinoCoreServerImpl) LoadSketch(ctx context.Context, req *rpc.LoadSke
 
 // Compile FIXMEDOC
 func (s *ArduinoCoreServerImpl) Compile(req *rpc.CompileRequest, stream rpc.ArduinoCoreService_CompileServer) error {
+	outStream, outCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.CompileResponse{OutStream: data}) })
+	errStream, errCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.CompileResponse{ErrStream: data}) })
 	resp, err := compile.Compile(
-		stream.Context(), req,
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.CompileResponse{OutStream: data}) }),
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.CompileResponse{ErrStream: data}) }),
+		stream.Context(), req, outStream, errStream,
 		func(p *rpc.TaskProgress) { stream.Send(&rpc.CompileResponse{Progress: p}) },
 		false) // Set debug to false
+	outStream.Close()
+	errStream.Close()
 	if err != nil {
 		return convertErrorToRPCStatus(err)
 	}
+	<-outCtx.Done()
+	<-errCtx.Done()
 	return stream.Send(resp)
 }
 
@@ -345,27 +349,31 @@ func (s *ArduinoCoreServerImpl) PlatformList(ctx context.Context, req *rpc.Platf
 
 // Upload FIXMEDOC
 func (s *ArduinoCoreServerImpl) Upload(req *rpc.UploadRequest, stream rpc.ArduinoCoreService_UploadServer) error {
-	resp, err := upload.Upload(
-		stream.Context(), req,
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadResponse{OutStream: data}) }),
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadResponse{ErrStream: data}) }),
-	)
+	outStream, outCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadResponse{OutStream: data}) })
+	errStream, errCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadResponse{ErrStream: data}) })
+	resp, err := upload.Upload(stream.Context(), req, outStream, errStream)
+	outStream.Close()
+	errStream.Close()
 	if err != nil {
 		return convertErrorToRPCStatus(err)
 	}
+	<-outCtx.Done()
+	<-errCtx.Done()
 	return stream.Send(resp)
 }
 
 // UploadUsingProgrammer FIXMEDOC
 func (s *ArduinoCoreServerImpl) UploadUsingProgrammer(req *rpc.UploadUsingProgrammerRequest, stream rpc.ArduinoCoreService_UploadUsingProgrammerServer) error {
-	resp, err := upload.UsingProgrammer(
-		stream.Context(), req,
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadUsingProgrammerResponse{OutStream: data}) }),
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadUsingProgrammerResponse{ErrStream: data}) }),
-	)
+	outStream, outCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadUsingProgrammerResponse{OutStream: data}) })
+	errStream, errCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.UploadUsingProgrammerResponse{ErrStream: data}) })
+	resp, err := upload.UsingProgrammer(stream.Context(), req, outStream, errStream)
+	outStream.Close()
+	errStream.Close()
 	if err != nil {
 		return convertErrorToRPCStatus(err)
 	}
+	<-outCtx.Done()
+	<-errCtx.Done()
 	return stream.Send(resp)
 }
 
@@ -377,14 +385,16 @@ func (s *ArduinoCoreServerImpl) SupportedUserFields(ctx context.Context, req *rp
 
 // BurnBootloader FIXMEDOC
 func (s *ArduinoCoreServerImpl) BurnBootloader(req *rpc.BurnBootloaderRequest, stream rpc.ArduinoCoreService_BurnBootloaderServer) error {
-	resp, err := upload.BurnBootloader(
-		stream.Context(), req,
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.BurnBootloaderResponse{OutStream: data}) }),
-		utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.BurnBootloaderResponse{ErrStream: data}) }),
-	)
+	outStream, outCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.BurnBootloaderResponse{OutStream: data}) })
+	errStream, errCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&rpc.BurnBootloaderResponse{ErrStream: data}) })
+	resp, err := upload.BurnBootloader(stream.Context(), req, outStream, errStream)
+	outStream.Close()
+	errStream.Close()
 	if err != nil {
 		return convertErrorToRPCStatus(err)
 	}
+	<-outCtx.Done()
+	<-errCtx.Done()
 	return stream.Send(resp)
 }
 

commands/daemon/debug.go

@@ -50,6 +50,7 @@ func (s *DebugService) Debug(stream dbg.DebugService_DebugServer) error {
 	// Launch debug recipe attaching stdin and out to grpc streaming
 	signalChan := make(chan os.Signal)
 	defer close(signalChan)
+	outStream, outCtx := utils.FeedStreamTo(func(data []byte) { stream.Send(&dbg.DebugResponse{Data: data}) })
 	resp, debugErr := cmd.Debug(stream.Context(), req,
 		utils.ConsumeStreamFrom(func() ([]byte, error) {
 			command, err := stream.Recv()
@@ -58,13 +59,13 @@ func (s *DebugService) Debug(stream dbg.DebugService_DebugServer) error {
 			}
 			return command.GetData(), err
 		}),
-		utils.FeedStreamTo(func(data []byte) {
-			stream.Send(&dbg.DebugResponse{Data: data})
-		}),
+		outStream,
 		signalChan)
+	outStream.Close()
 	if debugErr != nil {
 		return debugErr
 	}
+	<-outCtx.Done()
 	return stream.Send(resp)
 }