fix: Race when writing to a closed pipe

kylecarbs · kylecarbs · commit cec2de3291b2 · 2022-05-31T14:51:57.000Z
This is such an intermittent race it's difficult to track,
but regardless this is an improvement to the code.
diff --git a/provisionersdk/serve.go b/provisionersdk/serve.go
@@ -31,9 +31,9 @@ func Serve(ctx context.Context, server proto.DRPCProvisionerServer, options *Ser
 	if options.Listener == nil {
 		config := yamux.DefaultConfig()
 		config.LogOutput = io.Discard
-		stdio, err := yamux.Server(readWriteCloser{
-			ReadCloser: os.Stdin,
-			Writer:     os.Stdout,
+		stdio, err := yamux.Server(&readWriteCloser{
+			ReadCloser:  os.Stdin,
+			WriteCloser: os.Stdout,
 		}, config)
 		if err != nil {
 			return xerrors.Errorf("create yamux: %w", err)
@@ -54,6 +54,9 @@ func Serve(ctx context.Context, server proto.DRPCProvisionerServer, options *Ser
 	// short-lived processes that can be executed concurrently.
 	err = srv.Serve(ctx, options.Listener)
 	if err != nil {
+		if errors.Is(err, io.EOF) {
+			return nil
+		}
 		if errors.Is(err, context.Canceled) {
 			return nil
 		}
diff --git a/provisionersdk/transport.go b/provisionersdk/transport.go
@@ -22,16 +22,16 @@ func TransportPipe() (*yamux.Session, *yamux.Session) {
 	yamuxConfig := yamux.DefaultConfig()
 	yamuxConfig.LogOutput = io.Discard
 	client, err := yamux.Client(&readWriteCloser{
-		ReadCloser: clientReader,
-		Writer:     serverWriter,
+		ReadCloser:  clientReader,
+		WriteCloser: serverWriter,
 	}, yamuxConfig)
 	if err != nil {
 		panic(err)
 	}
 
 	server, err := yamux.Server(&readWriteCloser{
-		ReadCloser: serverReader,
-		Writer:     clientWriter,
+		ReadCloser:  serverReader,
+		WriteCloser: clientWriter,
 	}, yamuxConfig)
 	if err != nil {
 		panic(err)
@@ -46,7 +46,15 @@ func Conn(session *yamux.Session) drpc.Conn {
 
 type readWriteCloser struct {
 	io.ReadCloser
-	io.Writer
+	io.WriteCloser
+}
+
+func (c *readWriteCloser) Close() error {
+	err := c.ReadCloser.Close()
+	if err != nil {
+		return err
+	}
+	return c.WriteCloser.Close()
 }
 
 // Allows concurrent requests on a single dRPC connection.
