Merge pull request #4642 from oasisprotocol/kostko/stable/22.1.x/backport-4640

[BACKPORT/22.1.x] go/worker/compute: Improve proposed transaction handling

Author: kostko

.changelog/4640.bugfix.md

@@ -0,0 +1 @@
+go/worker/compute: Improve proposed transaction handling

.changelog/4640.feature.md

@@ -0,0 +1 @@
+runtime: Make persistent transaction check state configurable

Cargo.lock

@@ -78,6 +78,17 @@ type TransactionPool interface {
 	// SubmitTxNoWait adds the transaction into the transaction pool and returns immediately.
 	SubmitTxNoWait(ctx context.Context, tx []byte, meta *TransactionMeta) error
 
+	// SubmitProposedBatch adds the given (possibly new) transaction batch into the current
+	// proposal queue.
+	SubmitProposedBatch(batch [][]byte)
+
+	// PromoteProposedBatch promotes the specified transactions that are already in the transaction
+	// pool into the current proposal queue.
+	PromoteProposedBatch(batch []hash.Hash)
+
+	// ClearProposedBatch clears the proposal queue.
+	ClearProposedBatch()
+
 	// RemoveTxBatch removes a transaction batch from the transaction pool.
 	RemoveTxBatch(txs []hash.Hash)
 
@@ -183,11 +194,13 @@ type txPool struct {
 	schedulerTicker   *time.Ticker
 	schedulerNotifier *pubsub.Broker
 
+	proposedTxsLock sync.Mutex
+	proposedTxs     map[hash.Hash]*transaction.CheckedTransaction
+
 	blockInfoLock    sync.Mutex
 	blockInfo        *BlockInfo
 	lastRecheckRound uint64
 
-	epoCh       *channels.RingChannel
 	republishCh *channels.RingChannel
 
 	// roundWeightLimits is guarded by schedulerLock.
@@ -237,7 +250,7 @@ func (t *txPool) submitTx(ctx context.Context, rawTx []byte, meta *TransactionMe
 	// Skip recently seen transactions.
 	txHash := hash.NewFromBytes(rawTx)
 	if _, seen := t.seenCache.Peek(txHash); seen && !meta.Recheck {
-		t.logger.Debug("ignoring already seen transaction", "tx", rawTx)
+		t.logger.Debug("ignoring already seen transaction", "tx_hash", txHash)
 		return fmt.Errorf("duplicate transaction")
 	}
 
@@ -251,11 +264,13 @@ func (t *txPool) submitTx(ctx context.Context, rawTx []byte, meta *TransactionMe
 	// Queue transaction for checks.
 	t.logger.Debug("queuing transaction for check",
 		"tx", rawTx,
+		"tx_hash", txHash,
 		"recheck", meta.Recheck,
 	)
 	if err := t.checkTxQueue.Add(tx); err != nil {
 		t.logger.Warn("unable to queue transaction",
 			"tx", rawTx,
+			"tx_hash", txHash,
 			"err", err,
 		)
 		return err
@@ -269,6 +284,47 @@ func (t *txPool) submitTx(ctx context.Context, rawTx []byte, meta *TransactionMe
 	return nil
 }
 
+func (t *txPool) SubmitProposedBatch(batch [][]byte) {
+	// Also ingest into the regular pool (may fail).
+	for _, rawTx := range batch {
+		_ = t.SubmitTxNoWait(context.Background(), rawTx, &TransactionMeta{Local: false})
+	}
+
+	t.proposedTxsLock.Lock()
+	defer t.proposedTxsLock.Unlock()
+
+	for _, rawTx := range batch {
+		tx := transaction.RawCheckedTransaction(rawTx)
+		t.proposedTxs[tx.Hash()] = tx
+	}
+}
+
+func (t *txPool) PromoteProposedBatch(batch []hash.Hash) {
+	txs, missingTxs := t.GetKnownBatch(batch)
+	if len(missingTxs) > 0 {
+		t.logger.Debug("promoted proposed batch contains missing transactions",
+			"missing_tx_count", len(missingTxs),
+		)
+	}
+
+	t.proposedTxsLock.Lock()
+	defer t.proposedTxsLock.Unlock()
+
+	for _, tx := range txs {
+		if tx == nil {
+			continue
+		}
+		t.proposedTxs[tx.Hash()] = tx
+	}
+}
+
+func (t *txPool) ClearProposedBatch() {
+	t.proposedTxsLock.Lock()
+	defer t.proposedTxsLock.Unlock()
+
+	t.proposedTxs = make(map[hash.Hash]*transaction.CheckedTransaction)
+}
+
 func (t *txPool) RemoveTxBatch(txs []hash.Hash) {
 	t.schedulerLock.Lock()
 	defer t.schedulerLock.Unlock()
@@ -303,7 +359,6 @@ func (t *txPool) GetPrioritizedBatch(offset *hash.Hash, limit uint32) []*transac
 
 func (t *txPool) GetKnownBatch(batch []hash.Hash) ([]*transaction.CheckedTransaction, map[hash.Hash]int) {
 	t.schedulerLock.Lock()
-	defer t.schedulerLock.Unlock()
 
 	if t.schedulerQueue == nil {
 		result := make([]*transaction.CheckedTransaction, 0, len(batch))
@@ -312,9 +367,28 @@ func (t *txPool) GetKnownBatch(batch []hash.Hash) ([]*transaction.CheckedTransac
 			result = append(result, nil)
 			missing[txHash] = index
 		}
+		t.schedulerLock.Unlock()
 		return result, missing
 	}
-	return t.schedulerQueue.GetKnownBatch(batch)
+
+	txs, missingTxs := t.schedulerQueue.GetKnownBatch(batch)
+	t.schedulerLock.Unlock()
+
+	// Also check the proposed transactions set.
+	t.proposedTxsLock.Lock()
+	defer t.proposedTxsLock.Unlock()
+
+	for txHash, index := range missingTxs {
+		tx, exists := t.proposedTxs[txHash]
+		if !exists {
+			continue
+		}
+
+		delete(missingTxs, txHash)
+		txs[index] = tx
+	}
+
+	return txs, missingTxs
 }
 
 func (t *txPool) ProcessBlock(bi *BlockInfo) error {
@@ -327,7 +401,6 @@ func (t *txPool) ProcessBlock(bi *BlockInfo) error {
 			return fmt.Errorf("failed to update scheduler: %w", err)
 		}
 
-		t.epoCh.In() <- struct{}{}
 		// Force recheck on epoch transitions.
 		t.recheckTxCh.In() <- struct{}{}
 	}
@@ -409,12 +482,13 @@ func (t *txPool) WakeupScheduler() {
 
 func (t *txPool) Clear() {
 	t.schedulerLock.Lock()
-	defer t.schedulerLock.Unlock()
-
 	if t.schedulerQueue != nil {
 		t.schedulerQueue.Clear()
 	}
+	t.schedulerLock.Unlock()
+
 	t.seenCache.Clear()
+	t.ClearProposedBatch()
 
 	pendingScheduleSize.With(t.getMetricLabels()).Set(0)
 }
@@ -527,6 +601,7 @@ func (t *txPool) checkTxBatch(ctx context.Context, rr host.RichRuntime) {
 		if !res.IsSuccess() {
 			t.logger.Debug("check tx failed",
 				"tx", batch[i].Tx,
+				"tx_hash", batch[i].TxHash,
 				"result", res,
 				"recheck", batch[i].Meta.Recheck,
 			)
@@ -829,6 +904,10 @@ func New(
 		return nil, fmt.Errorf("error creating seen cache: %w", err)
 	}
 
+	// The transaction check queue should be 10% larger than the transaction pool to allow for some
+	// buffer in case the schedule queue is full and is being rechecked.
+	maxCheckTxQueueSize := (110 * cfg.MaxPoolSize) / 100
+
 	return &txPool{
 		logger:            logging.GetLogger("runtime/txpool"),
 		stopCh:            make(chan struct{}),
@@ -839,13 +918,13 @@ func New(
 		host:              host,
 		txPublisher:       txPublisher,
 		seenCache:         seenCache,
-		checkTxQueue:      newCheckTxQueue(cfg.MaxPoolSize, cfg.MaxCheckTxBatchSize),
+		checkTxQueue:      newCheckTxQueue(maxCheckTxQueueSize, cfg.MaxCheckTxBatchSize),
 		checkTxCh:         channels.NewRingChannel(1),
 		checkTxNotifier:   pubsub.NewBroker(false),
 		recheckTxCh:       channels.NewRingChannel(1),
 		schedulerTicker:   time.NewTicker(1 * time.Hour),
 		schedulerNotifier: pubsub.NewBroker(false),
-		epoCh:             channels.NewRingChannel(1),
+		proposedTxs:       make(map[hash.Hash]*transaction.CheckedTransaction),
 		republishCh:       channels.NewRingChannel(1),
 		roundWeightLimits: make(map[transaction.Weight]uint64),
 	}, nil

go/runtime/txpool/txpool.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
+	"github.com/eapache/channels"
 	core "github.com/libp2p/go-libp2p-core"
 	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/protocol"
@@ -146,6 +147,27 @@ func WithValidationFn(fn ValidationFunc) CallOption {
 	}
 }
 
+// AggregateFunc returns a result aggregation function.
+//
+// The function is passed the response and PeerFeedback instance. If the function returns true, the
+// client will continue to call other peers. If it returns false, processing will stop.
+type AggregateFunc func(rsp interface{}, pf PeerFeedback) bool
+
+// CallMultiOptions are per-multicall options
+type CallMultiOptions struct {
+	aggregateFn AggregateFunc
+}
+
+// CallMultiOption is a per-multicall option setter.
+type CallMultiOption func(opts *CallMultiOptions)
+
+// WithAggregateFn configures the response aggregation function to use.
+func WithAggregateFn(fn AggregateFunc) CallMultiOption {
+	return func(opts *CallMultiOptions) {
+		opts.aggregateFn = fn
+	}
+}
+
 // Client is an RPC client for a given protocol.
 type Client interface {
 	PeerManager
@@ -174,6 +196,7 @@ type Client interface {
 		body, rspTyp interface{},
 		maxPeerResponseTime time.Duration,
 		maxParallelRequests uint,
+		opts ...CallMultiOption,
 	) ([]interface{}, []PeerFeedback, error)
 }
 
@@ -276,9 +299,15 @@ func (c *client) CallMulti(
 	body, rspTyp interface{},
 	maxPeerResponseTime time.Duration,
 	maxParallelRequests uint,
+	opts ...CallMultiOption,
 ) ([]interface{}, []PeerFeedback, error) {
 	c.logger.Debug("call multiple", "method", method)
 
+	var co CallMultiOptions
+	for _, opt := range opts {
+		opt(&co)
+	}
+
 	// Prepare the request.
 	request := Request{
 		Method: method,
@@ -290,48 +319,75 @@ func (c *client) CallMulti(
 	pool.Resize(maxParallelRequests)
 	defer pool.Stop()
 
+	// Create a subcontext so we abort further requests if we are done early.
+	peerCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	// Requests results from peers.
 	type result struct {
 		rsp interface{}
 		pf  PeerFeedback
 		err error
 	}
-	var resultCh []chan *result
+	var resultChs []channels.SimpleOutChannel
 	for _, peer := range c.GetBestPeers() {
 		if !c.isPeerAcceptable(peer) {
 			continue
 		}
 
-		ch := make(chan *result, 1)
-		resultCh = append(resultCh, ch)
+		ch := channels.NewNativeChannel(channels.BufferCap(1))
+		resultChs = append(resultChs, ch)
 
 		pool.Submit(func() {
+			defer close(ch)
+
+			// Abort early in case we are done.
+			select {
+			case <-peerCtx.Done():
+				return
+			default:
+			}
+
 			rsp := reflect.New(reflect.TypeOf(rspTyp)).Interface()
-			pf, err := c.call(ctx, peer, &request, rsp, maxPeerResponseTime)
-			ch <- &result{rsp, pf, err}
-			close(ch)
+			pf, err := c.call(peerCtx, peer, &request, rsp, maxPeerResponseTime)
+			ch.In() <- &result{rsp, pf, err}
 		})
 	}
 
+	if len(resultChs) == 0 {
+		return nil, nil, nil
+	}
+	resultCh := channels.NewNativeChannel(channels.None)
+	channels.Multiplex(resultCh, resultChs...)
+
 	// Gather results.
 	var (
 		rsps []interface{}
 		pfs  []PeerFeedback
 	)
-	for _, ch := range resultCh {
-		select {
-		case <-ctx.Done():
-			return nil, nil, ctx.Err()
-		case result := <-ch:
-			// Ignore failed results.
-			if result.err != nil {
-				continue
-			}
+	for r := range resultCh.Out() {
+		result := r.(*result)
+
+		// Ignore failed results.
+		if result.err != nil {
+			continue
+		}
 
-			rsps = append(rsps, result.rsp)
-			pfs = append(pfs, result.pf)
+		rsps = append(rsps, result.rsp)
+		pfs = append(pfs, result.pf)
+
+		if co.aggregateFn != nil {
+			if !co.aggregateFn(result.rsp, result.pf) {
+				break
+			}
 		}
 	}
+
+	c.logger.Debug("received responses from peers",
+		"method", method,
+		"num_peers", len(rsps),
+	)
+
 	return rsps, pfs, nil
 }