From 25c2772d5fd908077927228597b084a1f1a0e66d Mon Sep 17 00:00:00 2001
From: chzyer <0@0xdf.com>
Date: Thu, 6 Oct 2016 12:35:07 +0800
Subject: [PATCH] fix deadlock in cancelable stdin (#89)

* fix deadlock in cancelable stdin

* fix notify back

* [cancelable stdin] fix c.notify
---
 std.go | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/std.go b/std.go
index ef0b6bb..e0c55ee 100644
--- a/std.go
+++ b/std.go
@@ -4,6 +4,7 @@ import (
 	"io"
 	"os"
 	"sync"
+	"sync/atomic"
 )
 
 var (
@@ -69,6 +70,7 @@ type CancelableStdin struct {
 	r      io.Reader
 	mutex  sync.Mutex
 	stop   chan struct{}
+	closed int32
 	notify chan struct{}
 	data   []byte
 	read   int
@@ -91,7 +93,11 @@ loop:
 		select {
 		case <-c.notify:
 			c.read, c.err = c.r.Read(c.data)
-			c.notify <- struct{}{}
+			select {
+			case c.notify <- struct{}{}:
+			case <-c.stop:
+				break loop
+			}
 		case <-c.stop:
 			break loop
 		}
@@ -101,9 +107,16 @@ loop:
 func (c *CancelableStdin) Read(b []byte) (n int, err error) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
+	if atomic.LoadInt32(&c.closed) == 1 {
+		return 0, io.EOF
+	}
 
 	c.data = b
-	c.notify <- struct{}{}
+	select {
+	case c.notify <- struct{}{}:
+	case <-c.stop:
+		return 0, io.EOF
+	}
 	select {
 	case <-c.notify:
 		return c.read, c.err
@@ -113,6 +126,8 @@ func (c *CancelableStdin) Read(b []byte) (n int, err error) {
 }
 
 func (c *CancelableStdin) Close() error {
-	close(c.stop)
+	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
+		close(c.stop)
+	}
 	return nil
 }