Fix exchanger deadlock during partial signature exchange (#4590)

Author: KaloyanTanev

dkg/exchanger.go

@@ -40,10 +40,22 @@ const (
 // sigTypeStore is a shorthand for a map of sigType to map of core.PubKey to slice of core.ParSignedData.
 type sigTypeStore map[sigType]map[core.PubKey][]core.ParSignedData
 
-// dataByPubkey maps a sigType to its map of public key to slice of core.ParSignedData..
+// dataByPubkey holds the partial signatures collected per sigType and the pending exchange queries
+// awaiting them. Both are guarded by lock.
 type dataByPubkey struct {
-	store sigTypeStore
-	lock  sync.Mutex
+	store   sigTypeStore
+	queries []exchangeQuery
+	lock    sync.Mutex
+}
+
+// exchangeQuery is a pending exchange call awaiting all its expected partial signatures for a sigType.
+type exchangeQuery struct {
+	sigType  sigType
+	expected int
+	// response is buffered (size 1) so resolveQueriesUnsafe never blocks, even when it runs on the
+	// exchange goroutine itself (pushPsigs may resolve queries synchronously via StoreInternal).
+	response chan<- map[core.PubKey][]core.ParSignedData
+	cancel   <-chan struct{}
 }
 
 // exchanger is responsible for exchanging partial signatures between peers on libp2p.
@@ -53,7 +65,6 @@ type exchanger struct {
 	sigTypes      map[sigType]bool
 	sigData       dataByPubkey
 	dutyGaterFunc func(duty core.Duty) bool
-	sigDatasChan  chan map[core.PubKey][]core.ParSignedData
 }
 
 func newExchanger(p2pNode host.Host, peerIdx int, peers []peer.ID, sigTypes []sigType, timeout time.Duration) *exchanger {
@@ -90,7 +101,6 @@ func newExchanger(p2pNode host.Host, peerIdx int, peers []peer.ID, sigTypes []si
 			lock:  sync.Mutex{},
 		},
 		dutyGaterFunc: dutyGaterFunc,
-		sigDatasChan:  make(chan map[core.PubKey][]core.ParSignedData, 1),
 	}
 
 	// Wiring core workflow components
@@ -112,60 +122,87 @@ func (e *exchanger) exchange(ctx context.Context, sigType sigType, set core.ParS
 		return nil, err
 	}
 
-	expectedSigs := len(set)
-
-	for {
-		select {
-		case sigDatas, ok := <-e.sigDatasChan:
-			if !ok {
-				return nil, errors.New("sigdata channel has been closed")
-			}
-
-			if len(sigDatas) == expectedSigs {
-				// We are done when we have ParSignedData of all the DVs from all each peer
-				return sigDatas, nil
-			}
-		case <-ctx.Done():
-			return nil, ctx.Err()
+	cancel := make(chan struct{})
+	defer close(cancel)
+
+	response := make(chan map[core.PubKey][]core.ParSignedData, 1)
+
+	// Register the query, then resolve immediately in case all expected signatures are already
+	// collected (e.g. peers delivered theirs before the StoreInternal above completed the threshold).
+	e.sigData.lock.Lock()
+	e.sigData.queries = append(e.sigData.queries, exchangeQuery{
+		sigType:  sigType,
+		expected: len(set),
+		response: response,
+		cancel:   cancel,
+	})
+	e.resolveQueriesUnsafe()
+	e.sigData.lock.Unlock()
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case data := <-response:
+		return data, nil
+	}
+}
+
+// resolveQueriesUnsafe delivers the collected signatures to every pending query whose sigType has
+// reached its expected count, dropping cancelled queries. It must be called with sigData.lock held.
+func (e *exchanger) resolveQueriesUnsafe() {
+	var remaining []exchangeQuery
+
+	for _, q := range e.sigData.queries {
+		if cancelled(q.cancel) {
+			continue
+		}
+
+		data := e.sigData.store[q.sigType]
+		if len(data) != q.expected {
+			remaining = append(remaining, q)
+			continue
 		}
+
+		// We are done when we have ParSignedData of all the DVs from each peer.
+		ret := make(map[core.PubKey][]core.ParSignedData, len(data))
+		maps.Copy(ret, data)
+
+		q.response <- ret // Never blocks: response is buffered and each query is resolved at most once.
+	}
+
+	e.sigData.queries = remaining
+}
+
+// cancelled returns true if the channel is closed.
+func cancelled(ch <-chan struct{}) bool {
+	select {
+	case <-ch:
+		return true
+	default:
+		return false
 	}
 }
 
-// pushPsigs is responsible for writing partial signature data to sigChan obtained from other peers.
-func (e *exchanger) pushPsigs(ctx context.Context, duty core.Duty, set map[core.PubKey][]core.ParSignedData) error {
+// pushPsigs stores partial signature data obtained from peers and resolves any pending query in
+// exchange whose expected signatures are now all collected.
+func (e *exchanger) pushPsigs(_ context.Context, duty core.Duty, set map[core.PubKey][]core.ParSignedData) error {
 	sigType := sigType(duty.Slot)
 
 	if !e.dutyGaterFunc(duty) {
 		return errors.New("unrecognized sigType", z.Int("sigType", int(sigType)))
 	}
 
 	e.sigData.lock.Lock()
+	defer e.sigData.lock.Unlock()
 
-	for pk, psigs := range set {
-		_, ok := e.sigData.store[sigType]
-		if !ok {
-			e.sigData.store[sigType] = map[core.PubKey][]core.ParSignedData{}
-		}
-
-		e.sigData.store[sigType][pk] = psigs
-	}
-
-	data, ok := e.sigData.store[sigType]
+	_, ok := e.sigData.store[sigType]
 	if !ok {
-		e.sigData.lock.Unlock()
-		return nil
+		e.sigData.store[sigType] = map[core.PubKey][]core.ParSignedData{}
 	}
 
-	ret := make(map[core.PubKey][]core.ParSignedData)
-	maps.Copy(ret, data)
-
-	e.sigData.lock.Unlock()
+	maps.Copy(e.sigData.store[sigType], set)
 
-	select {
-	case e.sigDatasChan <- ret:
-	case <-ctx.Done():
-		return errors.Wrap(ctx.Err(), "feed collected sig data")
-	}
+	e.resolveQueriesUnsafe()
 
 	return nil
 }

dkg/exchanger_internal_test.go

@@ -179,3 +179,48 @@ func TestExchanger(t *testing.T) {
 	// require that we encountered all the sigTypes expected
 	require.Len(t, actual, len(expectedSigTypes))
 }
+
+// TestExchangerPushPsigsNeverBlocks fires pushPsigs repeatedly with no exchange call draining
+// results, which must not block. A previous implementation used a shared size-1 channel that
+// deadlocked once full, especially when pushPsigs ran synchronously on the exchange goroutine.
+func TestExchangerPushPsigsNeverBlocks(t *testing.T) {
+	ctx := context.Background()
+
+	h := testutil.CreateHost(t, testutil.AvailableAddr(t))
+	peers := []peer.ID{h.ID()}
+
+	ex := newExchanger(h, 0, peers, []sigType{sigLock}, time.Second)
+
+	duty := core.NewSignatureDuty(uint64(sigLock))
+
+	newSet := func() map[core.PubKey][]core.ParSignedData {
+		return map[core.PubKey][]core.ParSignedData{
+			testutil.RandomCorePubKey(t): {core.NewPartialSignature(testutil.RandomCoreSignature(), 1)},
+		}
+	}
+
+	// Fire the threshold subscriber several times without any exchange call consuming results.
+	const iterations = 5
+
+	errs := make(chan error, 1)
+	done := make(chan struct{})
+
+	go func() {
+		for range iterations {
+			if err := ex.pushPsigs(ctx, duty, newSet()); err != nil {
+				errs <- err
+				return
+			}
+		}
+
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case err := <-errs:
+		require.NoError(t, err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("pushPsigs deadlocked")
+	}
+}