chore: thread context through to SendCustomMessage

The only way to unblock SendCustomMessage is if the peer activates,
disconnects or the server shuts down. This means that if the context is
cancelled, we will still wait until one of those other events happen.

With this commit we thread the context through to SendCustomMessage, so
that if the context is cancelled, we can return early. This improves the
cancellation semantics.

Author: gijswijs

rpcserver.go

@@ -9251,7 +9251,7 @@ func (r *rpcServer) RegisterRPCMiddleware(
 }
 
 // SendCustomMessage sends a custom peer message.
-func (r *rpcServer) SendCustomMessage(_ context.Context,
+func (r *rpcServer) SendCustomMessage(ctx context.Context,
 	req *lnrpc.SendCustomMessageRequest) (*lnrpc.SendCustomMessageResponse,
 	error) {
 
@@ -9261,7 +9261,7 @@ func (r *rpcServer) SendCustomMessage(_ context.Context,
 	}
 
 	err = r.server.SendCustomMessage(
-		peer, lnwire.MessageType(req.Type), req.Data,
+		ctx, peer, lnwire.MessageType(req.Type), req.Data,
 	)
 	switch {
 	case errors.Is(err, ErrPeerNotConnected):

server.go

@@ -5291,21 +5291,24 @@ func (s *server) applyChannelUpdate(update *lnwire.ChannelUpdate1,
 
 // SendCustomMessage sends a custom message to the peer with the specified
 // pubkey.
-func (s *server) SendCustomMessage(peerPub [33]byte, msgType lnwire.MessageType,
-	data []byte) error {
+func (s *server) SendCustomMessage(ctx context.Context, peerPub [33]byte,
+	msgType lnwire.MessageType, data []byte) error {
 
 	peer, err := s.FindPeerByPubStr(string(peerPub[:]))
 	if err != nil {
 		return err
 	}
 
-	// We'll wait until the peer is active.
+	// We'll wait until the peer is active, but also listen for
+	// cancellation.
 	select {
 	case <-peer.ActiveSignal():
 	case <-peer.QuitSignal():
 		return fmt.Errorf("peer %x disconnected", peerPub)
 	case <-s.quit:
 		return ErrServerShuttingDown
+	case <-ctx.Done():
+		return ctx.Err()
 	}
 
 	msg, err := lnwire.NewCustom(msgType, data)