Unblock handshake ch on close

Author: edaniels

association.go

@@ -1375,8 +1375,9 @@ func (a *Association) handleCookieEcho(c *chunkCookieEcho) []*packet {
 		a.storedCookieEcho = nil
 
 		a.setState(established)
-		// Note: This is a future place where the user could be notified (COMMUNICATION UP)
-		a.handshakeCompletedCh <- nil
+		if !a.completeHandshake(nil) {
+			return nil
+		}
 	}
 
 	p := &packet{
@@ -1404,8 +1405,7 @@ func (a *Association) handleCookieAck() {
 	a.storedCookieEcho = nil
 
 	a.setState(established)
-	// Note: This is a future place where the user could be notified (COMMUNICATION UP)
-	a.handshakeCompletedCh <- nil
+	a.completeHandshake(nil)
 }
 
 // The caller should hold the lock.
@@ -2698,13 +2698,13 @@ func (a *Association) onRetransmissionFailure(id int) {
 
 	if id == timerT1Init {
 		a.log.Errorf("[%s] retransmission failure: T1-init", a.name)
-		a.handshakeCompletedCh <- ErrHandshakeInitAck
+		a.completeHandshake(ErrHandshakeInitAck)
 		return
 	}
 
 	if id == timerT1Cookie {
 		a.log.Errorf("[%s] retransmission failure: T1-cookie", a.name)
-		a.handshakeCompletedCh <- ErrHandshakeCookieEcho
+		a.completeHandshake(ErrHandshakeCookieEcho)
 		return
 	}
 
@@ -2752,3 +2752,17 @@ func (a *Association) MaxMessageSize() uint32 {
 func (a *Association) SetMaxMessageSize(maxMsgSize uint32) {
 	atomic.StoreUint32(&a.maxMessageSize, maxMsgSize)
 }
+
+// completeHandshake sends the given error to  handshakeCompletedCh unless the read/write
+// side of the association closes before that can happen. It returns whether it was able
+// to send on the channel or not.
+func (a *Association) completeHandshake(handshakeErr error) bool {
+	select {
+	// Note: This is a future place where the user could be notified (COMMUNICATION UP)
+	case a.handshakeCompletedCh <- handshakeErr:
+		return true
+	case <-a.closeWriteLoopCh: // check the read/write sides for closure
+	case <-a.readLoopCloseCh:
+	}
+	return false
+}

association_test.go

@@ -3222,6 +3222,10 @@ func TestAssociation_Abort(t *testing.T) {
 
 // TestAssociation_createClientWithContext tests that the client is closed when the context is canceled.
 func TestAssociation_createClientWithContext(t *testing.T) {
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 5)
+	defer lim.Stop()
+
 	checkGoroutineLeaks(t)
 
 	udp1, udp2 := createUDPConnPair()