mempool: don't record a txid as seen until it has been fetched

common/mempool.go

@@ -122,7 +122,6 @@ func refreshMempoolTxns(ctx context.Context) error {
 		}
 
 		// We haven't fetched this transaction already.
-		g_txidSeen[txid(txidstr)] = struct{}{}
 		txidJSON, err := json.Marshal(txidstr)
 		if err != nil {
 			return err
@@ -131,9 +130,19 @@ func refreshMempoolTxns(ctx context.Context) error {
 		params := []json.RawMessage{txidJSON, json.RawMessage("1")}
 		result, rpcErr := RawRequest(ctx, "getrawtransaction", params)
 		if rpcErr != nil {
-			// Not an error; mempool transactions can disappear
+			// g_txidSeen is shared by every client streaming the mempool,
+			// but ctx belongs to the one whose goroutine is refreshing. If
+			// that client has gone away, stop without recording anything,
+			// so the next refresh fetches the rest for everyone else.
+			if err := ctx.Err(); err != nil {
+				return err
+			}
+			// Otherwise the transaction has most likely left the mempool
+			// and won't be listed again. Don't mark it seen: if the failure
+			// was transient, the next refresh retries it.
 			continue
 		}
+		g_txidSeen[txid(txidstr)] = struct{}{}
 
 		rawtx, err := ParseRawTransaction(result)
 		if err != nil {

common/mempool_refresh_test.go

@@ -0,0 +1,110 @@
+// Copyright (c) 2019-present The Zcash developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or https://www.opensource.org/licenses/mit-license.php .
+package common
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+)
+
+// g_txidSeen and g_txList are shared by every GetMempoolStream client, and
+// refreshMempoolTxns runs on whichever client's goroutine gets there first,
+// with that client's context. What one refresh records must hold for all.
+
+// mempoolRPC answers getrawmempool with txids and getrawtransaction with a
+// transaction, unless fail returns an error for that txid.
+func mempoolRPC(txids []string, fail func(ctx context.Context, txid string) error) func(context.Context, string, []json.RawMessage) (json.RawMessage, error) {
+	return func(ctx context.Context, method string, params []json.RawMessage) (json.RawMessage, error) {
+		if err := ctx.Err(); err != nil {
+			// The real client aborts the HTTP request once ctx is done.
+			return nil, err
+		}
+		switch method {
+		case "getrawmempool":
+			return json.Marshal(txids)
+		case "getrawtransaction":
+			var id string
+			json.Unmarshal(params[0], &id)
+			if err := fail(ctx, id); err != nil {
+				return nil, err
+			}
+			return json.Marshal(map[string]string{"hex": "aabb"})
+		}
+		return nil, errors.New("unexpected method " + method)
+	}
+}
+
+func TestMempoolRefreshAfterClientCancelFetchesRemainingTxns(t *testing.T) {
+	defer resetGlobals()
+	txids := []string{"tx-a", "tx-b", "tx-c"}
+
+	// The client whose goroutine is refreshing disconnects after tx-a.
+	ctx, cancel := context.WithCancel(context.Background())
+	RawRequest = mempoolRPC(txids, func(_ context.Context, id string) error {
+		if id == "tx-a" {
+			cancel()
+		}
+		return nil
+	})
+	if err := refreshMempoolTxns(ctx); !errors.Is(err, context.Canceled) {
+		t.Errorf("refresh with a cancelled client returned %v, want context.Canceled", err)
+	}
+
+	// The next refresh, on behalf of another client, must still fetch the
+	// transactions the cancelled one never got to.
+	RawRequest = mempoolRPC(txids, func(context.Context, string) error { return nil })
+	if err := refreshMempoolTxns(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if len(g_txList) != len(txids) {
+		t.Fatalf("mempool list has %d transactions, want %d", len(g_txList), len(txids))
+	}
+}
+
+func TestMempoolRefreshRetriesTransientFetchFailure(t *testing.T) {
+	defer resetGlobals()
+	txids := []string{"tx-a", "tx-b", "tx-c"}
+
+	failed := false
+	RawRequest = mempoolRPC(txids, func(_ context.Context, id string) error {
+		if id == "tx-b" && !failed {
+			failed = true
+			return errors.New("connection reset by peer")
+		}
+		return nil
+	})
+	if err := refreshMempoolTxns(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if len(g_txList) != 2 {
+		t.Fatalf("first refresh has %d transactions, want 2", len(g_txList))
+	}
+	// tx-b is still in the mempool, so the next refresh fetches it.
+	if err := refreshMempoolTxns(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if len(g_txList) != 3 {
+		t.Fatalf("second refresh has %d transactions, want 3", len(g_txList))
+	}
+}
+
+func TestMempoolRefreshDoesNotRefetchKnownTxns(t *testing.T) {
+	defer resetGlobals()
+	txids := []string{"tx-a", "tx-b"}
+	fetches := 0
+	RawRequest = mempoolRPC(txids, func(context.Context, string) error {
+		fetches++
+		return nil
+	})
+	for i := 0; i < 3; i++ {
+		if err := refreshMempoolTxns(context.Background()); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if fetches != 2 || len(g_txList) != 2 {
+		t.Fatalf("fetched %d times for %d transactions, want 2 and 2", fetches, len(g_txList))
+	}
+}