diff --git a/dkg/exchanger.go b/dkg/exchanger.go index 2e7f99efd..04ca2f865 100644 --- a/dkg/exchanger.go +++ b/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,27 +122,70 @@ 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) { @@ -140,32 +193,16 @@ func (e *exchanger) pushPsigs(ctx context.Context, duty core.Duty, set map[core. } 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 } diff --git a/dkg/exchanger_internal_test.go b/dkg/exchanger_internal_test.go index 7d7acfff7..b2ce77fcb 100644 --- a/dkg/exchanger_internal_test.go +++ b/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") + } +}