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 }