Forráskód Böngészése

Fix race around read and write deadlines in Stream (#52)

Filip Ochnik 7 éve
szülő
commit
683f49123a
1 módosított fájl, 12 hozzáadás és 8 törlés
  1. 12 8
      stream.go

+ 12 - 8
stream.go

@@ -47,8 +47,8 @@ type Stream struct {
 	recvNotifyCh chan struct{}
 	sendNotifyCh chan struct{}
 
-	readDeadline  time.Time
-	writeDeadline time.Time
+	readDeadline  atomic.Value // time.Time
+	writeDeadline atomic.Value // time.Time
 }
 
 // newStream is used to construct a new stream within
@@ -67,6 +67,8 @@ func newStream(session *Session, id uint32, state streamState) *Stream {
 		recvNotifyCh: make(chan struct{}, 1),
 		sendNotifyCh: make(chan struct{}, 1),
 	}
+	s.readDeadline.Store(time.Time{})
+	s.writeDeadline.Store(time.Time{})
 	return s
 }
 
@@ -122,8 +124,9 @@ START:
 WAIT:
 	var timeout <-chan time.Time
 	var timer *time.Timer
-	if !s.readDeadline.IsZero() {
-		delay := s.readDeadline.Sub(time.Now())
+	readDeadline := s.readDeadline.Load().(time.Time)
+	if !readDeadline.IsZero() {
+		delay := readDeadline.Sub(time.Now())
 		timer = time.NewTimer(delay)
 		timeout = timer.C
 	}
@@ -200,8 +203,9 @@ START:
 
 WAIT:
 	var timeout <-chan time.Time
-	if !s.writeDeadline.IsZero() {
-		delay := s.writeDeadline.Sub(time.Now())
+	writeDeadline := s.writeDeadline.Load().(time.Time)
+	if !writeDeadline.IsZero() {
+		delay := writeDeadline.Sub(time.Now())
 		timeout = time.After(delay)
 	}
 	select {
@@ -435,13 +439,13 @@ func (s *Stream) SetDeadline(t time.Time) error {
 
 // SetReadDeadline sets the deadline for future Read calls.
 func (s *Stream) SetReadDeadline(t time.Time) error {
-	s.readDeadline = t
+	s.readDeadline.Store(t)
 	return nil
 }
 
 // SetWriteDeadline sets the deadline for future Write calls
 func (s *Stream) SetWriteDeadline(t time.Time) error {
-	s.writeDeadline = t
+	s.writeDeadline.Store(t)
 	return nil
 }