authmailbox: wait for multisub reconnects in restart test

- Wait for the shared MultiSubscription to report active subscriptions
  again after the mock server restarts before publishing msg2.
- The test already waited for the two direct clients, but not for the
  multi-subscription, which left a race where the post-restart message
  could be published before that consumer had reconnected.
- Because this subscription uses an empty filter, a missed publish is
  not recovered from backlog delivery and the test times out on
  readMultiSub.
- Keep the fix test-only so mailbox runtime behavior is unchanged.
- This makes TestServerClientAuthAndRestart pass reliably under race
  runs and matches the failure seen in CI shard 0.

Author: bhandras

authmailbox/client_test.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/btcsuite/btclog/v2"
+	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/fn"
 	"github.com/lightninglabs/taproot-assets/internal/test"
 	"github.com/lightninglabs/taproot-assets/proof"
@@ -169,6 +170,32 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	t.Cleanup(func() {
 		require.NoError(t, multiSub.Stop())
 	})
+	assertMultiSubConnected := func(targetKeys ...keychain.KeyDescriptor) {
+		t.Helper()
+
+		serverURL := url.URL{Host: clientCfg.ServerAddress}
+		require.Eventually(t, func() bool {
+			multiSub.RLock()
+			defer multiSub.RUnlock()
+
+			client, ok := multiSub.clients[serverURL]
+			if !ok {
+				return false
+			}
+
+			for _, targetKey := range targetKeys {
+				key := asset.ToSerialized(targetKey.PubKey)
+				subscription, ok := client.subscriptions[key]
+				if !ok || !subscription.IsSubscribed() {
+					return false
+				}
+			}
+
+			return true
+		}, testTimeout, testMinBackoff)
+	}
+	assertMultiSubConnected(clientKey1, clientKey2)
+
 	msgChan := multiSub.MessageChan()
 	readMultiSub := func(targetID ...uint64) {
 		t.Helper()
@@ -221,6 +248,7 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	harness.Start(t)
 	client1.assertConnected(t)
 	client2.assertConnected(t)
+	assertMultiSubConnected(clientKey1, clientKey2)
 
 	// Let's send another message to all clients.
 	msg2 := &Message{

authmailbox/receive_subscription.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"math"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/btcsuite/btclog/v2"
@@ -61,6 +62,7 @@ type receiveSubscription struct {
 	serverStream clientStream
 	streamMutex  sync.RWMutex
 	streamCancel func()
+	subscribed   atomic.Bool
 
 	authOkChan chan struct{}
 	msgChan    chan<- *ReceivedMessages
@@ -198,10 +200,7 @@ func (s *receiveSubscription) wait(backoff time.Duration) error {
 // IsSubscribed returns true if at least one account is in an active state and
 // the subscription stream to the server was established successfully.
 func (s *receiveSubscription) IsSubscribed() bool {
-	s.streamMutex.RLock()
-	defer s.streamMutex.RUnlock()
-
-	return s.serverStream != nil
+	return s.subscribed.Load()
 }
 
 // connectServerStream opens the initial connection to the server for the stream
@@ -412,6 +411,8 @@ func (s *receiveSubscription) readIncomingStream(ctx context.Context) {
 		// The server confirms the account subscription. Nothing for us
 		// to do here.
 		case *respTypeAuthSuccess:
+			s.subscribed.Store(true)
+
 			// Inform the subscription about the arrived auth
 			// confirmation.
 			select {
@@ -476,6 +477,7 @@ func (s *receiveSubscription) closeStream(ctx context.Context) error {
 
 	s.streamMutex.Lock()
 	defer s.streamMutex.Unlock()
+	s.subscribed.Store(false)
 
 	if s.streamCancel != nil {
 		s.streamCancel()