Merge pull request #127 from lattwood/lattwood/handle_shutdown_edgecase

Streams should check for Session shutdown when waiting for data & clean up timers

stream.go

@@ -143,14 +143,17 @@ WAIT:
 		timeout = timer.C
 	}
 	select {
+	case <-s.session.shutdownCh:
 	case <-s.recvNotifyCh:
-		if timer != nil {
-			timer.Stop()
-		}
-		goto START
 	case <-timeout:
 		return 0, ErrTimeout
 	}
+	if timer != nil {
+		if !timer.Stop() {
+			<-timeout
+		}
+	}
+	goto START
 }
 
 // Write is used to write to the stream
@@ -219,17 +222,25 @@ START:
 
 WAIT:
 	var timeout <-chan time.Time
+	var timer *time.Timer
 	writeDeadline := s.writeDeadline.Load().(time.Time)
 	if !writeDeadline.IsZero() {
 		delay := time.Until(writeDeadline)
-		timeout = time.After(delay)
+		timer = time.NewTimer(delay)
+		timeout = timer.C
 	}
 	select {
+	case <-s.session.shutdownCh:
 	case <-s.sendNotifyCh:
-		goto START
 	case <-timeout:
 		return 0, ErrTimeout
 	}
+	if timer != nil {
+		if !timer.Stop() {
+			<-timeout
+		}
+	}
+	goto START
 }
 
 // sendFlags determines any flags that are appropriate