diff --git a/session_test.go b/session_test.go index 00156ed..0b53430 100644 --- a/session_test.go +++ b/session_test.go @@ -10,6 +10,7 @@ import ( "net" "net/http" _ "net/http/pprof" + "strings" "sync" "testing" "time" @@ -471,6 +472,34 @@ func TestRandomFrame(t *testing.T) { cli.Close() } +func TestReadDeadline(t *testing.T) { + cli, err := net.Dial("tcp", "127.0.0.1:19999") + if err != nil { + t.Fatal(err) + } + session, _ := Client(cli, nil) + stream, _ := session.OpenStream() + const N = 100 + buf := make([]byte, 10) + var readErr error + for i := 0; i < N; i++ { + msg := fmt.Sprintf("hello%v", i) + stream.Write([]byte(msg)) + stream.SetReadDeadline(time.Now().Add(-1 * time.Minute)) + if _, readErr = stream.Read(buf); readErr != nil { + break + } + } + if readErr != nil { + if !strings.Contains(readErr.Error(), "i/o timeout") { + t.Fatalf("Wrong error: %v", readErr) + } + } else { + t.Fatal("No error when reading with past deadline") + } + session.Close() +} + func BenchmarkAcceptClose(b *testing.B) { cli, err := net.Dial("tcp", "127.0.0.1:19999") if err != nil { diff --git a/stream.go b/stream.go index f7e28c7..dfd7b04 100644 --- a/stream.go +++ b/stream.go @@ -4,21 +4,23 @@ import ( "bytes" "sync" "sync/atomic" + "time" "github.com/pkg/errors" ) // Stream implements io.ReadWriteCloser type Stream struct { - id uint32 - rstflag int32 - sess *Session - buffer bytes.Buffer - bufferLock sync.Mutex - frameSize int - chReadEvent chan struct{} // notify a read event - die chan struct{} // flag the stream has closed - dieLock sync.Mutex + id uint32 + rstflag int32 + sess *Session + buffer bytes.Buffer + bufferLock sync.Mutex + frameSize int + chReadEvent chan struct{} // notify a read event + die chan struct{} // flag the stream has closed + dieLock sync.Mutex + readDeadline int64 } // newStream initiates a Stream struct @@ -34,10 +36,20 @@ func newStream(id uint32, frameSize int, sess *Session) *Stream { // Read implements io.ReadWriteCloser func (s *Stream) Read(b []byte) (n int, err error) { + var deadline <-chan time.Time + d := atomic.LoadInt64(&s.readDeadline) + if d > 0 { + timer := time.NewTimer(time.Duration(d - time.Now().UnixNano())) + defer timer.Stop() + deadline = timer.C + } + READ: select { case <-s.die: return 0, errors.New(errBrokenPipe) + case <-deadline: + return n, errTimeout default: } @@ -56,6 +68,8 @@ READ: select { case <-s.chReadEvent: goto READ + case <-deadline: + return n, errTimeout case <-s.die: return 0, errors.New(errBrokenPipe) } @@ -94,6 +108,12 @@ func (s *Stream) Close() error { } } +// SetReadDeadline sets the read deadline as defined by +// net.Conn.SetReadDeadline. +func (s *Stream) SetReadDeadline(t time.Time) { + atomic.StoreInt64(&s.readDeadline, t.UnixNano()) +} + // session closes the stream func (s *Stream) sessionClose() { s.dieLock.Lock() @@ -151,3 +171,11 @@ func (s *Stream) notifyReadEvent() { func (s *Stream) markRST() { atomic.StoreInt32(&s.rstflag, 1) } + +var errTimeout error = &timeoutError{} + +type timeoutError struct{} + +func (e *timeoutError) Error() string { return "i/o timeout" } +func (e *timeoutError) Timeout() bool { return true } +func (e *timeoutError) Temporary() bool { return true }