fix a potential logic race

xtaci · Apr 18, 2019 · d568d23 · d568d23
1 parent 11b388c
commit d568d23
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/session.go b/session.go
@@ -113,9 +113,14 @@ func (s *Session) OpenStream() (*Stream, error) {
 	}
 
 	s.streamLock.Lock()
-	s.streams[sid] = stream
-	s.streamLock.Unlock()
-	return stream, nil
+	defer s.streamLock.Unlock()
+	select {
+	case <-s.die:
+		return nil, errors.New(errBrokenPipe)
+	default:
+		s.streams[sid] = stream
+		return stream, nil
+	}
 }
 
 // AcceptStream is used to block until the next available stream