Wire per-topic partial-messages SubOpts flags (libp2p#444)

lucassaldanha · lucassaldanha · commit 5971b2e2f5a7 · 2026-04-21T17:10:57.000+12:00
Step 1 of the partial-messages extension: plumb SubOpts.requestsPartial /
SubOpts.supportsSendingPartial through subscribe announcements in both
directions, and track the per-peer-per-topic receive state.

- AbstractRouter parses the flags with the spec-mandated coercion
  (supportsSendingPartial := requestsPartial || supportsSendingPartial)
  and zeroes them on subscribe=false.
- New enqueueSubscribe hook unifies outbound subscribe enqueueing so
  GossipRouter can attach per-topic flags in a single override.
- GossipRouter exposes setTopicPartialFlags(topic, ...) to configure
  flags advertised for a locally-subscribed topic, and stores inbound
  flags in a new PartialSubscriptionState (plain HashMap on the pubsub
  event loop). State is cleaned on peer disconnect, topic unsubscribe,
  and per-peer unsubscribe.
- Outbound unsubscribe MUST NOT carry partial flags; enforced at the
  SubscriptionPart wire-build site.

No routing behaviour changes yet. See docs/partial-messages.md §4.5,
§5, §6.1 for context.
diff --git a/libp2p/src/main/kotlin/io/libp2p/pubsub/AbstractRouter.kt b/libp2p/src/main/kotlin/io/libp2p/pubsub/AbstractRouter.kt
@@ -147,11 +147,22 @@ abstract class AbstractRouter(
     override fun onPeerActive(peer: PeerHandler) {
         val partsQueue = pendingRpcParts.getQueue(peer)
         subscribedTopics.forEach {
-            partsQueue.addSubscribe(it)
+            enqueueSubscribe(partsQueue, it)
         }
         flushPending(peer)
     }
 
+    /**
+     * Enqueues a subscribe announcement for [topic] onto [partsQueue].
+     *
+     * The default implementation emits a bare subscribe with no per-topic options.
+     * Subclasses (e.g. GossipRouter) override this to attach per-topic options
+     * such as partial-message flags.
+     */
+    protected open fun enqueueSubscribe(partsQueue: RpcPartsQueue, topic: Topic) {
+        partsQueue.addSubscribe(topic)
+    }
+
     protected open fun notifyMalformedMessage(peer: PeerHandler) {}
     protected open fun notifyUnseenMessage(peer: PeerHandler, msg: PubsubMessage) {}
     protected open fun notifyNonSubscribedMessage(peer: PeerHandler, msg: Rpc.Message) {}
@@ -172,7 +183,16 @@ abstract class AbstractRouter(
         }
 
         try {
-            val subscriptions = msg.subscriptionsList.map { PubsubSubscription(it.topicid, it.subscribe) }
+            val subscriptions = msg.subscriptionsList.map {
+                // Per partial-messages spec: flags MUST be ignored on subscribe=false, and the
+                // receiving side coerces supportsSendingPartial := requestsPartial || supportsSendingPartial.
+                PubsubSubscription(
+                    topic = it.topicid,
+                    subscribe = it.subscribe,
+                    requestsPartial = it.subscribe && it.requestsPartial,
+                    supportsSendingPartial = it.subscribe && (it.supportsSendingPartial || it.requestsPartial)
+                )
+            }
             subscriptionFilter
                 .filterIncomingSubscriptions(subscriptions, peersTopics.getByFirst(peer))
                 .forEach { handleMessageSubscriptions(peer, it) }
@@ -301,7 +321,7 @@ abstract class AbstractRouter(
         }
     }
 
-    private fun handleMessageSubscriptions(peer: PeerHandler, msg: PubsubSubscription) {
+    protected open fun handleMessageSubscriptions(peer: PeerHandler, msg: PubsubSubscription) {
         if (msg.subscribe) {
             peersTopics.add(peer, msg.topic)
         } else {
@@ -319,7 +339,7 @@ abstract class AbstractRouter(
     }
 
     protected open fun subscribe(topic: Topic) {
-        activePeers.forEach { pendingRpcParts.getQueue(it).addSubscribe(topic) }
+        activePeers.forEach { enqueueSubscribe(pendingRpcParts.getQueue(it), topic) }
         subscribedTopics += topic
     }
 
diff --git a/libp2p/src/main/kotlin/io/libp2p/pubsub/PubsubRouter.kt b/libp2p/src/main/kotlin/io/libp2p/pubsub/PubsubRouter.kt
@@ -13,7 +13,12 @@ typealias Topic = String
 typealias MessageId = WBytes
 typealias PubsubMessageFactory = (Rpc.Message) -> PubsubMessage
 
-data class PubsubSubscription(val topic: Topic, val subscribe: Boolean)
+data class PubsubSubscription(
+    val topic: Topic,
+    val subscribe: Boolean,
+    val requestsPartial: Boolean = false,
+    val supportsSendingPartial: Boolean = false
+)
 
 interface PubsubMessage {
     val protobufMessage: Rpc.Message
diff --git a/libp2p/src/main/kotlin/io/libp2p/pubsub/RpcPartsQueue.kt b/libp2p/src/main/kotlin/io/libp2p/pubsub/RpcPartsQueue.kt
@@ -9,14 +9,27 @@ interface RpcPartsQueue {
     fun addPublish(message: Rpc.Message)
 
     fun addSubscribe(topic: Topic) {
-        addSubscription(topic, SubscriptionStatus.Subscribed)
+        addSubscribe(topic, requestsPartial = false, supportsSendingPartial = false)
+    }
+
+    fun addSubscribe(topic: Topic, requestsPartial: Boolean, supportsSendingPartial: Boolean) {
+        addSubscription(topic, SubscriptionStatus.Subscribed, requestsPartial, supportsSendingPartial)
     }
 
     fun addUnsubscribe(topic: Topic) {
-        addSubscription(topic, SubscriptionStatus.Unsubscribed)
+        addSubscription(topic, SubscriptionStatus.Unsubscribed, requestsPartial = false, supportsSendingPartial = false)
+    }
+
+    fun addSubscription(topic: Topic, status: SubscriptionStatus) {
+        addSubscription(topic, status, requestsPartial = false, supportsSendingPartial = false)
     }
 
-    fun addSubscription(topic: Topic, status: SubscriptionStatus)
+    fun addSubscription(
+        topic: Topic,
+        status: SubscriptionStatus,
+        requestsPartial: Boolean,
+        supportsSendingPartial: Boolean
+    )
 
     fun takeMerged(): List<Rpc.RPC>
 }
@@ -38,11 +51,20 @@ open class DefaultRpcPartsQueue : RpcPartsQueue {
         }
     }
 
-    protected data class SubscriptionPart(val topic: Topic, val status: RpcPartsQueue.SubscriptionStatus) : AbstractPart {
+    protected data class SubscriptionPart(
+        val topic: Topic,
+        val status: RpcPartsQueue.SubscriptionStatus,
+        val requestsPartial: Boolean = false,
+        val supportsSendingPartial: Boolean = false
+    ) : AbstractPart {
         override fun appendToBuilder(builder: Rpc.RPC.Builder) {
-            builder.addSubscriptionsBuilder().apply {
-                setTopicid(topic)
-                setSubscribe(status == RpcPartsQueue.SubscriptionStatus.Subscribed)
+            val subBuilder = builder.addSubscriptionsBuilder()
+            subBuilder.topicid = topic
+            subBuilder.subscribe = status == RpcPartsQueue.SubscriptionStatus.Subscribed
+            // Per spec: partial flags MUST NOT be sent on unsubscribe (subscribe=false).
+            if (status == RpcPartsQueue.SubscriptionStatus.Subscribed) {
+                if (requestsPartial) subBuilder.requestsPartial = true
+                if (supportsSendingPartial) subBuilder.supportsSendingPartial = true
             }
         }
     }
@@ -57,8 +79,13 @@ open class DefaultRpcPartsQueue : RpcPartsQueue {
         addPart(PublishPart(message))
     }
 
-    override fun addSubscription(topic: Topic, status: RpcPartsQueue.SubscriptionStatus) {
-        addPart(SubscriptionPart(topic, status))
+    override fun addSubscription(
+        topic: Topic,
+        status: RpcPartsQueue.SubscriptionStatus,
+        requestsPartial: Boolean,
+        supportsSendingPartial: Boolean
+    ) {
+        addPart(SubscriptionPart(topic, status, requestsPartial, supportsSendingPartial))
     }
 
     override fun takeMerged(): List<Rpc.RPC> {
diff --git a/libp2p/src/main/kotlin/io/libp2p/pubsub/gossip/GossipRouter.kt b/libp2p/src/main/kotlin/io/libp2p/pubsub/gossip/GossipRouter.kt
@@ -134,6 +134,35 @@ open class GossipRouter(
     override val pendingRpcParts = PendingRpcPartsMap<GossipRpcPartsQueue> { DefaultGossipRpcPartsQueue(params) }
 
     val gossipExtensionsState = GossipExtensionsState(gossipExtensionsConfig)
+    val partialSubscriptionState = PartialSubscriptionState()
+
+    /**
+     * Local per-topic subscription options that affect outbound subscribe announcements.
+     * Accessed only on the pubsub event loop.
+     */
+    private val localTopicPartialFlags: MutableMap<Topic, PartialSubFlags> = mutableMapOf()
+
+    /**
+     * Configures the partial-messages flags advertised on this node's subscribe
+     * announcements for [topic]. Must be called before [subscribe] for the flags
+     * to take effect on the initial announcement; a subsequent call will affect
+     * later re-announcements (e.g. on new peer activation).
+     *
+     * Per spec, the send-side also applies the coercion
+     * `supportsSendingPartial := requestsPartial || supportsSendingPartial`.
+     */
+    fun setTopicPartialFlags(topic: Topic, requestsPartial: Boolean, supportsSendingPartial: Boolean) {
+        runOnEventThread {
+            if (!requestsPartial && !supportsSendingPartial) {
+                localTopicPartialFlags -= topic
+            } else {
+                localTopicPartialFlags[topic] = PartialSubFlags(
+                    requestsPartial = requestsPartial,
+                    supportsSendingPartial = requestsPartial || supportsSendingPartial
+                )
+            }
+        }
+    }
 
     private fun setBackOff(peer: PeerHandler, topic: Topic) = setBackOff(peer, topic, params.pruneBackoff.toMillis())
     private fun setBackOff(peer: PeerHandler, topic: Topic, delay: Long) {
@@ -161,9 +190,28 @@ open class GossipRouter(
         acceptRequestsWhitelist -= peer
         pendingRpcParts.popQueue(peer) // discard them
         gossipExtensionsState.onPeerDisconnected(peer.peerId)
+        partialSubscriptionState.onPeerDisconnected(peer.peerId)
         super.onPeerDisconnected(peer)
     }
 
+    override fun enqueueSubscribe(partsQueue: RpcPartsQueue, topic: Topic) {
+        val flags = localTopicPartialFlags[topic] ?: PartialSubFlags.NONE
+        partsQueue.addSubscribe(topic, flags.requestsPartial, flags.supportsSendingPartial)
+    }
+
+    override fun handleMessageSubscriptions(peer: PeerHandler, msg: PubsubSubscription) {
+        super.handleMessageSubscriptions(peer, msg)
+        if (msg.subscribe) {
+            partialSubscriptionState.setPeerFlags(
+                msg.topic,
+                peer.peerId,
+                PartialSubFlags(msg.requestsPartial, msg.supportsSendingPartial)
+            )
+        } else {
+            partialSubscriptionState.removePeerFlags(msg.topic, peer.peerId)
+        }
+    }
+
     override fun onPeerActive(peer: PeerHandler) {
         super.onPeerActive(peer)
         eventBroadcaster.notifyConnected(peer.peerId, peer.getRemoteAddress())
@@ -615,6 +663,8 @@ open class GossipRouter(
         super.unsubscribe(topic)
         mesh[topic]?.copy()?.forEach { prune(it, topic) }
         mesh -= topic
+        localTopicPartialFlags -= topic
+        partialSubscriptionState.removeTopic(topic)
     }
 
     private fun catchingHeartbeat() {
diff --git a/libp2p/src/main/kotlin/io/libp2p/pubsub/gossip/PartialSubscriptionState.kt b/libp2p/src/main/kotlin/io/libp2p/pubsub/gossip/PartialSubscriptionState.kt
@@ -0,0 +1,66 @@
+package io.libp2p.pubsub.gossip
+
+import io.libp2p.core.PeerId
+import io.libp2p.pubsub.Topic
+
+data class PartialSubFlags(
+    val requestsPartial: Boolean,
+    val supportsSendingPartial: Boolean
+) {
+    companion object {
+        val NONE = PartialSubFlags(requestsPartial = false, supportsSendingPartial = false)
+    }
+}
+
+/**
+ * Per-topic, per-peer partial-messages subscription state.
+ *
+ * Tracks, for each `(topic, peer)`, the remote peer's `requestsPartial` /
+ * `supportsSendingPartial` flags as most recently announced via a subscribe
+ * `SubOpts`. Unsubscribes and peer disconnects drop the corresponding state.
+ *
+ * NOT thread-safe: accessed only on the pubsub event loop.
+ */
+class PartialSubscriptionState {
+
+    private val byTopic: MutableMap<Topic, MutableMap<PeerId, PartialSubFlags>> = mutableMapOf()
+
+    fun setPeerFlags(topic: Topic, peer: PeerId, flags: PartialSubFlags) {
+        if (flags == PartialSubFlags.NONE) {
+            removePeerFlags(topic, peer)
+            return
+        }
+        byTopic.getOrPut(topic) { mutableMapOf() }[peer] = flags
+    }
+
+    fun removePeerFlags(topic: Topic, peer: PeerId) {
+        val peers = byTopic[topic] ?: return
+        peers.remove(peer)
+        if (peers.isEmpty()) byTopic.remove(topic)
+    }
+
+    fun removeTopic(topic: Topic) {
+        byTopic.remove(topic)
+    }
+
+    fun onPeerDisconnected(peer: PeerId) {
+        val emptied = mutableListOf<Topic>()
+        for ((topic, peers) in byTopic) {
+            peers.remove(peer)
+            if (peers.isEmpty()) emptied += topic
+        }
+        emptied.forEach { byTopic.remove(it) }
+    }
+
+    fun peerFlags(topic: Topic, peer: PeerId): PartialSubFlags =
+        byTopic[topic]?.get(peer) ?: PartialSubFlags.NONE
+
+    fun peerRequestsPartial(topic: Topic, peer: PeerId) =
+        peerFlags(topic, peer).requestsPartial
+
+    fun peerSupportsSendingPartial(topic: Topic, peer: PeerId) =
+        peerFlags(topic, peer).supportsSendingPartial
+
+    internal fun snapshot(): Map<Topic, Map<PeerId, PartialSubFlags>> =
+        byTopic.mapValues { (_, v) -> v.toMap() }
+}
diff --git a/libp2p/src/test/kotlin/io/libp2p/pubsub/gossip/PartialSubscriptionStateTest.kt b/libp2p/src/test/kotlin/io/libp2p/pubsub/gossip/PartialSubscriptionStateTest.kt
diff --git a/libp2p/src/test/kotlin/io/libp2p/pubsub/gossip/extensions/PartialSubscriptionWireTest.kt b/libp2p/src/test/kotlin/io/libp2p/pubsub/gossip/extensions/PartialSubscriptionWireTest.kt
