Compare commits

...

3 Commits

Author SHA1 Message Date
chzyer 5197e2b2ba [cancelable stdin] fix c.notify 2016-10-06 12:33:21 +08:00
chzyer 17b536725b fix notify back 2016-10-06 12:31:47 +08:00
chzyer fa8d922787 fix deadlock in cancelable stdin 2016-10-06 12:30:22 +08:00
1 changed files with 18 additions and 3 deletions

19
std.go
View File

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