Don't send output chunk if channel is closed

Fixes puppetlabs-toy-chest#562

Signed-off-by: Enis Inan <[email protected]>

plugin/outputStream.go

@@ -13,7 +13,7 @@ type OutputStream struct {
 	id         ExecPacketType
 	ch         chan ExecOutputChunk
 	closer     *multiCloser
-	mux        sync.Mutex
+	mux        sync.RWMutex
 	closed     bool
 }
 
@@ -27,6 +27,13 @@ func (s *OutputStream) sendError(timestamp time.Time, err error) {
 
 // WriteWithTimestamp writes the given data with the specified timestamp
 func (s *OutputStream) WriteWithTimestamp(timestamp time.Time, data []byte) error {
+	s.mux.RLock()
+	defer s.mux.RUnlock()
+
+	if s.closed {
+		return nil
+	}
+
 	select {
 	case <-s.ctx.Done():
 		s.sendError(timestamp, s.ctx.Err())
@@ -79,4 +86,4 @@ func (c *multiCloser) Close() {
 		close(c.ch)
 	}
 	c.mux.Unlock()
-}
+}

plugin/outputStream_test.go

@@ -21,10 +21,10 @@ func newOutputStream(ctx context.Context) *OutputStream {
 	ch := make(chan ExecOutputChunk, 1)
 	return &OutputStream{
 		ctx: ctx,
-		id: Stdout,
-		ch: ch,
+		id:  Stdout,
+		ch:  ch,
 		closer: &multiCloser{
-			ch: ch,
+			ch:        ch,
 			countdown: 1,
 		},
 	}
@@ -66,7 +66,7 @@ func (suite *OutputStreamTestSuite) TestWrite() {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
 	stream := newOutputStream(ctx)
-	
+
 	// Test a successful write
 	data := []byte("data")
 	nw, writeErr := stream.Write(data)
@@ -187,6 +187,21 @@ func (suite *OutputStreamTestSuite) TestCloseWithError_ConsecutiveCloses() {
 	suite.assertClosed(stream)
 }
 
+func (suite *OutputStreamTestSuite) TestCloseWithError_AllSubsequentSendsNoop() {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	stream := newOutputStream(ctx)
+	stream.CloseWithError(nil)
+	suite.assertClosed(stream)
+
+	// Note that the test will panic if stream.Write sends something on
+	// the closed channel.
+	_, err := stream.Write([]byte("foo"))
+	suite.NoError(err)
+	cancelFunc()
+	_, err = stream.Write([]byte("bar"))
+	suite.NoError(err)
+}
+
 func TestOutputStream(t *testing.T) {
 	suite.Run(t, new(OutputStreamTestSuite))
 }