mirror of
https://github.com/Chia-Network/chia-blockchain.git
synced 2024-10-04 06:48:03 +03:00
Refactor subscriptions and add remove functionality (#17249)
* Refactor PeerSubscriptions * Rename ph to puzzle in subscriptions * Add subscription list test * Fix test * Improve code quality and logging * Remove redundant line
This commit is contained in:
Rigidity
GitHub
parent
714c4d1985
commit
0cafc83194
@ -1154,7 +1154,7 @@ class FullNode:
|
||||
hints_to_add, _ = get_hints_and_subscription_coin_ids(
|
||||
state_change_summary,
|
||||
self.subscriptions.has_coin_subscription,
|
||||
self.subscriptions.has_ph_subscription,
|
||||
self.subscriptions.has_puzzle_subscription,
|
||||
)
|
||||
await self.hint_store.add_hints(hints_to_add)
|
||||
# Note that end_height is not necessarily the peak at this
|
||||
@ -1484,7 +1484,7 @@ class FullNode:
|
||||
hints_to_add, lookup_coin_ids = get_hints_and_subscription_coin_ids(
|
||||
state_change_summary,
|
||||
self.subscriptions.has_coin_subscription,
|
||||
self.subscriptions.has_ph_subscription,
|
||||
self.subscriptions.has_puzzle_subscription,
|
||||
)
|
||||
await self.hint_store.add_hints(hints_to_add)
|
||||
|
||||
|
||||
@ -1497,7 +1497,7 @@ class FullNodeAPI:
|
||||
# the returned puzzle hashes are the ones we ended up subscribing to.
|
||||
# It will have filtered duplicates and ones exceeding the subscription
|
||||
# limit.
|
||||
puzzle_hashes = self.full_node.subscriptions.add_ph_subscriptions(
|
||||
puzzle_hashes = self.full_node.subscriptions.add_puzzle_subscriptions(
|
||||
peer.peer_node_id, request.puzzle_hashes, max_subscriptions
|
||||
)
|
||||
|
||||
|
||||
@ -9,7 +9,7 @@ from chia.types.blockchain_format.sized_bytes import bytes32
|
||||
def get_hints_and_subscription_coin_ids(
|
||||
state_change_summary: StateChangeSummary,
|
||||
has_coin_subscription: Callable[[bytes32], bool],
|
||||
has_ph_subscription: Callable[[bytes32], bool],
|
||||
has_puzzle_subscription: Callable[[bytes32], bool],
|
||||
) -> Tuple[List[Tuple[bytes32, bytes]], List[bytes32]]:
|
||||
# Precondition: all hints passed in are max 32 bytes long
|
||||
# Returns the hints that we need to add to the DB, and the coin ids that need to be looked up
|
||||
@ -26,7 +26,7 @@ def get_hints_and_subscription_coin_ids(
|
||||
lookup_coin_ids.add(coin_id)
|
||||
|
||||
def add_if_ph_subscription(puzzle_hash: bytes32, coin_id: bytes32) -> None:
|
||||
if has_ph_subscription(puzzle_hash):
|
||||
if has_puzzle_subscription(puzzle_hash):
|
||||
lookup_coin_ids.add(coin_id)
|
||||
|
||||
for spend_id, puzzle_hash in state_change_summary.removals:
|
||||
|
||||
@ -9,136 +9,196 @@ from chia.types.blockchain_format.sized_bytes import bytes32
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# The PeerSubscriptions class is essentially a multi-index container. It can be
|
||||
# indexed by peer_id, coin_id and puzzle_hash.
|
||||
@dataclass(frozen=True)
|
||||
class PeerSubscriptions:
|
||||
# TODO: use NewType all over to describe these various uses of the same types
|
||||
# Puzzle Hash : Set[Peer ID]
|
||||
_coin_subscriptions: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
# Puzzle Hash : Set[Peer ID]
|
||||
_ph_subscriptions: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
# Peer ID: Set[Coin ids]
|
||||
_peer_coin_ids: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
# Peer ID: Set[puzzle_hash]
|
||||
_peer_puzzle_hash: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
# Peer ID: subscription count
|
||||
_peer_sub_counter: Dict[bytes32, int] = field(default_factory=dict, init=False)
|
||||
class SubscriptionSet:
|
||||
_subscriptions_for_peer: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
_peers_for_subscription: Dict[bytes32, Set[bytes32]] = field(default_factory=dict, init=False)
|
||||
|
||||
def has_ph_subscription(self, ph: bytes32) -> bool:
|
||||
return ph in self._ph_subscriptions
|
||||
def add_subscription(self, peer_id: bytes32, item: bytes32) -> bool:
|
||||
peers = self._peers_for_subscription.setdefault(item, set())
|
||||
|
||||
def has_coin_subscription(self, coin_id: bytes32) -> bool:
|
||||
return coin_id in self._coin_subscriptions
|
||||
if peer_id in peers:
|
||||
return False
|
||||
|
||||
def add_ph_subscriptions(self, peer_id: bytes32, phs: List[bytes32], max_items: int) -> Set[bytes32]:
|
||||
"""
|
||||
returns the puzzle hashes that were actually subscribed to. These may be
|
||||
fewer than requested in case:
|
||||
* there are duplicate puzzle_hashes
|
||||
* some puzzle hashes are already subscribed to
|
||||
* the max_items limit is exceeded
|
||||
"""
|
||||
subscriptions = self._subscriptions_for_peer.setdefault(peer_id, set())
|
||||
subscriptions.add(item)
|
||||
peers.add(peer_id)
|
||||
|
||||
puzzle_hash_peers = self._peer_puzzle_hash.setdefault(peer_id, set())
|
||||
existing_sub_count = self._peer_sub_counter.setdefault(peer_id, 0)
|
||||
return True
|
||||
|
||||
ret: Set[bytes32] = set()
|
||||
def remove_subscription(self, peer_id: bytes32, item: bytes32) -> bool:
|
||||
subscriptions = self._subscriptions_for_peer.get(peer_id)
|
||||
|
||||
# if we've reached the limit on number of subscriptions, just bail
|
||||
if existing_sub_count >= max_items:
|
||||
log.info(
|
||||
"peer_id: %s reached max number of puzzle-hash subscriptions. "
|
||||
"Not all its coin states will be reported",
|
||||
peer_id,
|
||||
)
|
||||
return ret
|
||||
if subscriptions is None or item not in subscriptions:
|
||||
return False
|
||||
|
||||
# decrement this counter as we go, to know if we've hit the limit of
|
||||
# number of subscriptions
|
||||
subscriptions_left = max_items - existing_sub_count
|
||||
peers = self._peers_for_subscription[item]
|
||||
peers.remove(peer_id)
|
||||
subscriptions.remove(item)
|
||||
|
||||
for ph in phs:
|
||||
ph_sub = self._ph_subscriptions.setdefault(ph, set())
|
||||
if peer_id in ph_sub:
|
||||
continue
|
||||
if len(subscriptions) == 0:
|
||||
self._subscriptions_for_peer.pop(peer_id)
|
||||
|
||||
ret.add(ph)
|
||||
ph_sub.add(peer_id)
|
||||
puzzle_hash_peers.add(ph)
|
||||
self._peer_sub_counter[peer_id] += 1
|
||||
subscriptions_left -= 1
|
||||
if len(peers) == 0:
|
||||
self._peers_for_subscription.pop(item)
|
||||
|
||||
if subscriptions_left == 0:
|
||||
log.info(
|
||||
"peer_id: %s reached max number of puzzle-hash subscriptions. "
|
||||
"Not all its coin states will be reported",
|
||||
peer_id,
|
||||
)
|
||||
break
|
||||
return ret
|
||||
return True
|
||||
|
||||
def add_coin_subscriptions(self, peer_id: bytes32, coin_ids: List[bytes32], max_items: int) -> None:
|
||||
coin_id_peers = self._peer_coin_ids.setdefault(peer_id, set())
|
||||
existing_sub_count = self._peer_sub_counter.setdefault(peer_id, 0)
|
||||
def has_subscription(self, item: bytes32) -> bool:
|
||||
return item in self._peers_for_subscription
|
||||
|
||||
# if we've reached the limit on number of subscriptions, just bail
|
||||
if existing_sub_count >= max_items:
|
||||
log.info(
|
||||
"peer_id: %s reached max number of coin subscriptions. Not all its coin states will be reported",
|
||||
peer_id,
|
||||
)
|
||||
return
|
||||
|
||||
# decrement this counter as we go, to know if we've hit the limit of
|
||||
# number of subscriptions
|
||||
subscriptions_left = max_items - existing_sub_count
|
||||
|
||||
for coin_id in coin_ids:
|
||||
coin_sub = self._coin_subscriptions.setdefault(coin_id, set())
|
||||
if peer_id in coin_sub:
|
||||
continue
|
||||
|
||||
coin_sub.add(peer_id)
|
||||
coin_id_peers.add(coin_id)
|
||||
self._peer_sub_counter[peer_id] += 1
|
||||
subscriptions_left -= 1
|
||||
|
||||
if subscriptions_left == 0:
|
||||
log.info(
|
||||
"peer_id: %s reached max number of coin subscriptions. Not all its coin states will be reported",
|
||||
peer_id,
|
||||
)
|
||||
break
|
||||
def count_subscriptions(self, peer_id: bytes32) -> int:
|
||||
return len(self._subscriptions_for_peer.get(peer_id, {}))
|
||||
|
||||
def remove_peer(self, peer_id: bytes32) -> None:
|
||||
counter = 0
|
||||
puzzle_hashes = self._peer_puzzle_hash.get(peer_id)
|
||||
if puzzle_hashes is not None:
|
||||
for ph in puzzle_hashes:
|
||||
subs = self._ph_subscriptions[ph]
|
||||
subs.remove(peer_id)
|
||||
counter += 1
|
||||
if subs == set():
|
||||
self._ph_subscriptions.pop(ph)
|
||||
self._peer_puzzle_hash.pop(peer_id)
|
||||
for item in self._subscriptions_for_peer.pop(peer_id, {}):
|
||||
self._peers_for_subscription[item].remove(peer_id)
|
||||
|
||||
coin_ids = self._peer_coin_ids.get(peer_id)
|
||||
if coin_ids is not None:
|
||||
for coin_id in coin_ids:
|
||||
subs = self._coin_subscriptions[coin_id]
|
||||
subs.remove(peer_id)
|
||||
counter += 1
|
||||
if subs == set():
|
||||
self._coin_subscriptions.pop(coin_id)
|
||||
self._peer_coin_ids.pop(peer_id)
|
||||
if len(self._peers_for_subscription[item]) == 0:
|
||||
self._peers_for_subscription.pop(item)
|
||||
|
||||
if peer_id in self._peer_sub_counter:
|
||||
num_subs = self._peer_sub_counter.pop(peer_id)
|
||||
assert num_subs == counter
|
||||
def subscriptions(self, peer_id: bytes32) -> Set[bytes32]:
|
||||
return self._subscriptions_for_peer.get(peer_id, set())
|
||||
|
||||
def peers(self, item: bytes32) -> Set[bytes32]:
|
||||
return self._peers_for_subscription.get(item, set())
|
||||
|
||||
def total_count(self) -> int:
|
||||
return len(self._peers_for_subscription)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PeerSubscriptions:
|
||||
_puzzle_subscriptions: SubscriptionSet = field(default_factory=SubscriptionSet)
|
||||
_coin_subscriptions: SubscriptionSet = field(default_factory=SubscriptionSet)
|
||||
|
||||
def has_puzzle_subscription(self, puzzle_hash: bytes32) -> bool:
|
||||
return self._puzzle_subscriptions.has_subscription(puzzle_hash)
|
||||
|
||||
def has_coin_subscription(self, coin_id: bytes32) -> bool:
|
||||
return self._coin_subscriptions.has_subscription(coin_id)
|
||||
|
||||
def peer_subscription_count(self, peer_id: bytes32) -> int:
|
||||
puzzle_subscriptions = self._puzzle_subscriptions.count_subscriptions(peer_id)
|
||||
coin_subscriptions = self._coin_subscriptions.count_subscriptions(peer_id)
|
||||
return puzzle_subscriptions + coin_subscriptions
|
||||
|
||||
def add_puzzle_subscriptions(self, peer_id: bytes32, puzzle_hashes: List[bytes32], max_items: int) -> Set[bytes32]:
|
||||
"""
|
||||
Adds subscriptions until max_items is reached. Filters out duplicates and returns all additions.
|
||||
"""
|
||||
|
||||
subscription_count = self.peer_subscription_count(peer_id)
|
||||
added: Set[bytes32] = set()
|
||||
|
||||
def limit_reached() -> Set[bytes32]:
|
||||
log.info(
|
||||
"Peer %s attempted to exceed the subscription limit while adding puzzle subscriptions.",
|
||||
peer_id,
|
||||
)
|
||||
return added
|
||||
|
||||
# If the subscription limit is reached, bail.
|
||||
if subscription_count >= max_items:
|
||||
return limit_reached()
|
||||
|
||||
# Decrement this counter to know if we've hit the subscription limit.
|
||||
subscriptions_left = max_items - subscription_count
|
||||
|
||||
for puzzle_hash in puzzle_hashes:
|
||||
if not self._puzzle_subscriptions.add_subscription(peer_id, puzzle_hash):
|
||||
continue
|
||||
|
||||
subscriptions_left -= 1
|
||||
added.add(puzzle_hash)
|
||||
|
||||
if subscriptions_left == 0:
|
||||
return limit_reached()
|
||||
|
||||
return added
|
||||
|
||||
def add_coin_subscriptions(self, peer_id: bytes32, coin_ids: List[bytes32], max_items: int) -> Set[bytes32]:
|
||||
"""
|
||||
Adds subscriptions until max_items is reached. Filters out duplicates and returns all additions.
|
||||
"""
|
||||
|
||||
subscription_count = self.peer_subscription_count(peer_id)
|
||||
added: Set[bytes32] = set()
|
||||
|
||||
def limit_reached() -> Set[bytes32]:
|
||||
log.info(
|
||||
"Peer %s attempted to exceed the subscription limit while adding coin subscriptions.",
|
||||
peer_id,
|
||||
)
|
||||
return added
|
||||
|
||||
# If the subscription limit is reached, bail.
|
||||
if subscription_count >= max_items:
|
||||
return limit_reached()
|
||||
|
||||
# Decrement this counter to know if we've hit the subscription limit.
|
||||
subscriptions_left = max_items - subscription_count
|
||||
|
||||
for coin_id in coin_ids:
|
||||
if not self._coin_subscriptions.add_subscription(peer_id, coin_id):
|
||||
continue
|
||||
|
||||
subscriptions_left -= 1
|
||||
added.add(coin_id)
|
||||
|
||||
if subscriptions_left == 0:
|
||||
return limit_reached()
|
||||
|
||||
return added
|
||||
|
||||
def remove_puzzle_subscriptions(self, peer_id: bytes32, puzzle_hashes: List[bytes32]) -> Set[bytes32]:
|
||||
"""
|
||||
Removes subscriptions. Filters out duplicates and returns all removals.
|
||||
"""
|
||||
|
||||
removed: Set[bytes32] = set()
|
||||
|
||||
for puzzle_hash in puzzle_hashes:
|
||||
if not self._puzzle_subscriptions.remove_subscription(peer_id, puzzle_hash):
|
||||
continue
|
||||
|
||||
removed.add(puzzle_hash)
|
||||
|
||||
return removed
|
||||
|
||||
def remove_coin_subscriptions(self, peer_id: bytes32, coin_ids: List[bytes32]) -> Set[bytes32]:
|
||||
"""
|
||||
Removes subscriptions. Filters out duplicates and returns all removals.
|
||||
"""
|
||||
|
||||
removed: Set[bytes32] = set()
|
||||
|
||||
for coin_id in coin_ids:
|
||||
if not self._coin_subscriptions.remove_subscription(peer_id, coin_id):
|
||||
continue
|
||||
|
||||
removed.add(coin_id)
|
||||
|
||||
return removed
|
||||
|
||||
def remove_peer(self, peer_id: bytes32) -> None:
|
||||
self._puzzle_subscriptions.remove_peer(peer_id)
|
||||
self._coin_subscriptions.remove_peer(peer_id)
|
||||
|
||||
def coin_subscriptions(self, peer_id: bytes32) -> Set[bytes32]:
|
||||
return self._coin_subscriptions.subscriptions(peer_id)
|
||||
|
||||
def puzzle_subscriptions(self, peer_id: bytes32) -> Set[bytes32]:
|
||||
return self._puzzle_subscriptions.subscriptions(peer_id)
|
||||
|
||||
def peers_for_coin_id(self, coin_id: bytes32) -> Set[bytes32]:
|
||||
return self._coin_subscriptions.get(coin_id, set())
|
||||
return self._coin_subscriptions.peers(coin_id)
|
||||
|
||||
def peers_for_puzzle_hash(self, puzzle_hash: bytes32) -> Set[bytes32]:
|
||||
return self._ph_subscriptions.get(puzzle_hash, set())
|
||||
return self._puzzle_subscriptions.peers(puzzle_hash)
|
||||
|
||||
def coin_subscription_count(self) -> int:
|
||||
return self._coin_subscriptions.total_count()
|
||||
|
||||
def puzzle_subscription_count(self) -> int:
|
||||
return self._puzzle_subscriptions.total_count()
|
||||
|
||||
@ -20,21 +20,21 @@ ph4 = bytes32(b"h" * 32)
|
||||
def test_has_ph_sub() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1], 100)
|
||||
assert ret == {ph1}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1, ph2], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1, ph2], 100)
|
||||
# we have already subscribed to ph1, it's filtered in the returned list
|
||||
assert ret == {ph2}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
|
||||
# note that this is technically a type error as well.
|
||||
# we can remove these asserts once we have type checking
|
||||
@ -43,8 +43,8 @@ def test_has_ph_sub() -> None:
|
||||
|
||||
sub.remove_peer(peer1)
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
|
||||
|
||||
def test_has_coin_sub() -> None:
|
||||
@ -65,8 +65,8 @@ def test_has_coin_sub() -> None:
|
||||
|
||||
# note that this is technically a type error as well.
|
||||
# we can remove these asserts once we have type checking
|
||||
assert sub.has_ph_subscription(coin1) is False
|
||||
assert sub.has_ph_subscription(coin2) is False
|
||||
assert sub.has_puzzle_subscription(coin1) is False
|
||||
assert sub.has_puzzle_subscription(coin2) is False
|
||||
|
||||
sub.remove_peer(peer1)
|
||||
|
||||
@ -119,34 +119,34 @@ def test_overlapping_coin_subscriptions() -> None:
|
||||
def test_overlapping_ph_subscriptions() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == set()
|
||||
assert sub.peers_for_puzzle_hash(ph2) == set()
|
||||
|
||||
# subscribed to different phs
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1], 100)
|
||||
assert ret == {ph1}
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == set()
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer2, [ph2], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer2, [ph2], 100)
|
||||
assert ret == {ph2}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == {peer2}
|
||||
|
||||
# peer1 is now subscribing to both phs
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph2], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph2], 100)
|
||||
assert ret == {ph2}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == {peer1, peer2}
|
||||
@ -154,8 +154,8 @@ def test_overlapping_ph_subscriptions() -> None:
|
||||
# removing peer1 still leaves the subscription to ph2
|
||||
sub.remove_peer(peer1)
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == set()
|
||||
assert sub.peers_for_puzzle_hash(ph2) == {peer2}
|
||||
@ -164,19 +164,19 @@ def test_overlapping_ph_subscriptions() -> None:
|
||||
def test_ph_sub_limit() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1, ph2, ph3, ph4], 3)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1, ph2, ph3, ph4], 3)
|
||||
# we only ended up subscribing to 3 puzzle hashes because of the limit
|
||||
assert ret == {ph1, ph2, ph3}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_ph_subscription(ph3) is True
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph3) is True
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == {peer1}
|
||||
@ -184,10 +184,10 @@ def test_ph_sub_limit() -> None:
|
||||
assert sub.peers_for_puzzle_hash(ph4) == set()
|
||||
|
||||
# peer1 should still be limited
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph4], 3)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph4], 3)
|
||||
assert ret == set()
|
||||
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
assert sub.peers_for_puzzle_hash(ph4) == set()
|
||||
|
||||
# peer1 is also limied on coin subscriptions
|
||||
@ -197,10 +197,10 @@ def test_ph_sub_limit() -> None:
|
||||
assert sub.peers_for_coin_id(coin1) == set()
|
||||
|
||||
# peer2 is has its own limit
|
||||
ret = sub.add_ph_subscriptions(peer2, [ph4], 3)
|
||||
ret = sub.add_puzzle_subscriptions(peer2, [ph4], 3)
|
||||
assert ret == {ph4}
|
||||
|
||||
assert sub.has_ph_subscription(ph4) is True
|
||||
assert sub.has_puzzle_subscription(ph4) is True
|
||||
assert sub.peers_for_puzzle_hash(ph4) == {peer2}
|
||||
|
||||
sub.remove_peer(peer1)
|
||||
@ -210,18 +210,18 @@ def test_ph_sub_limit() -> None:
|
||||
def test_ph_sub_limit_incremental() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1], 2)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1], 2)
|
||||
assert ret == {ph1}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == set()
|
||||
@ -229,13 +229,13 @@ def test_ph_sub_limit_incremental() -> None:
|
||||
assert sub.peers_for_puzzle_hash(ph4) == set()
|
||||
|
||||
# this will cross the limit. Only ph2 will be added
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph2, ph3], 2)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph2, ph3], 2)
|
||||
assert ret == {ph2}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
assert sub.peers_for_puzzle_hash(ph1) == {peer1}
|
||||
assert sub.peers_for_puzzle_hash(ph2) == {peer1}
|
||||
@ -272,10 +272,10 @@ def test_coin_sub_limit() -> None:
|
||||
assert sub.peers_for_coin_id(coin4) == set()
|
||||
|
||||
# peer1 is also limied on ph subscriptions
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1], 3)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1], 3)
|
||||
assert ret == set()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.peers_for_puzzle_hash(ph1) == set()
|
||||
|
||||
# peer2 is has its own limit
|
||||
@ -327,31 +327,77 @@ def test_coin_sub_limit_incremental() -> None:
|
||||
def test_ph_subscription_duplicates() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1, ph2, ph3], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1, ph2, ph3], 100)
|
||||
assert ret == {ph1, ph2, ph3}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_ph_subscription(ph3) is True
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph3) is True
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
# only ph4 is new, the others are duplicates and ignored
|
||||
ret = sub.add_ph_subscriptions(peer1, [ph1, ph2, ph3, ph4], 100)
|
||||
ret = sub.add_puzzle_subscriptions(peer1, [ph1, ph2, ph3, ph4], 100)
|
||||
assert ret == {ph4}
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is True
|
||||
assert sub.has_ph_subscription(ph2) is True
|
||||
assert sub.has_ph_subscription(ph3) is True
|
||||
assert sub.has_ph_subscription(ph4) is True
|
||||
assert sub.has_puzzle_subscription(ph1) is True
|
||||
assert sub.has_puzzle_subscription(ph2) is True
|
||||
assert sub.has_puzzle_subscription(ph3) is True
|
||||
assert sub.has_puzzle_subscription(ph4) is True
|
||||
|
||||
sub.remove_peer(peer1)
|
||||
|
||||
assert sub.has_ph_subscription(ph1) is False
|
||||
assert sub.has_ph_subscription(ph2) is False
|
||||
assert sub.has_ph_subscription(ph3) is False
|
||||
assert sub.has_ph_subscription(ph4) is False
|
||||
assert sub.has_puzzle_subscription(ph1) is False
|
||||
assert sub.has_puzzle_subscription(ph2) is False
|
||||
assert sub.has_puzzle_subscription(ph3) is False
|
||||
assert sub.has_puzzle_subscription(ph4) is False
|
||||
|
||||
|
||||
def test_remove_ph_subscriptions() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
added = sub.add_puzzle_subscriptions(peer1, [ph1, ph2, ph3, ph4, ph4], 100)
|
||||
assert added == {ph1, ph2, ph3, ph4}
|
||||
|
||||
removed = sub.remove_puzzle_subscriptions(peer1, list(added))
|
||||
assert removed == added
|
||||
|
||||
# These have already been removed.
|
||||
assert len(sub.remove_puzzle_subscriptions(peer1, [ph1, ph2])) == 0
|
||||
|
||||
assert sub.peer_subscription_count(peer1) == 0
|
||||
|
||||
for ph in removed:
|
||||
assert not sub.has_puzzle_subscription(ph)
|
||||
|
||||
|
||||
def test_remove_coin_subscriptions() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
added = sub.add_coin_subscriptions(peer1, [coin1, coin2, coin3, coin4, coin4], 100)
|
||||
assert added == {coin1, coin2, coin3, coin4}
|
||||
|
||||
removed = sub.remove_coin_subscriptions(peer1, list(added))
|
||||
assert removed == added
|
||||
|
||||
# These have already been removed.
|
||||
assert len(sub.remove_coin_subscriptions(peer1, [coin1, coin2])) == 0
|
||||
|
||||
assert sub.peer_subscription_count(peer1) == 0
|
||||
|
||||
for coin_id in removed:
|
||||
assert not sub.has_coin_subscription(coin_id)
|
||||
|
||||
|
||||
def test_subscription_list() -> None:
|
||||
sub = PeerSubscriptions()
|
||||
|
||||
sub.add_coin_subscriptions(peer1, [coin1, coin2], 4)
|
||||
sub.add_puzzle_subscriptions(peer1, [ph1, ph2], 4)
|
||||
|
||||
assert sub.coin_subscriptions(peer1) == {coin1, coin2}
|
||||
assert sub.puzzle_subscriptions(peer1) == {ph1, ph2}
|
||||
|
||||
@ -962,7 +962,7 @@ async def test_cat_change_detection(
|
||||
await time_out_assert(20, check_wallets, 2, wallet_node_0)
|
||||
cat_wallet = wallet_node_0.wallet_state_manager.wallets[uint32(2)]
|
||||
await time_out_assert(20, cat_wallet.get_confirmed_balance, cat_amount_1)
|
||||
assert not full_node_api.full_node.subscriptions.has_ph_subscription(puzzlehash_unhardened)
|
||||
assert not full_node_api.full_node.subscriptions.has_puzzle_subscription(puzzlehash_unhardened)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
||||
@ -677,23 +677,23 @@ class TestSimpleSyncProtocol:
|
||||
msg_response = await full_node_api.register_interest_in_puzzle_hash(msg, con)
|
||||
assert msg_response.type == ProtocolMessageTypes.respond_to_ph_update.value
|
||||
s = full_node_api.full_node.subscriptions
|
||||
assert len(s._ph_subscriptions) == 2
|
||||
assert s.has_ph_subscription(phs[0])
|
||||
assert s.has_ph_subscription(phs[1])
|
||||
assert not s.has_ph_subscription(phs[2])
|
||||
assert not s.has_ph_subscription(phs[3])
|
||||
assert s.puzzle_subscription_count() == 2
|
||||
assert s.has_puzzle_subscription(phs[0])
|
||||
assert s.has_puzzle_subscription(phs[1])
|
||||
assert not s.has_puzzle_subscription(phs[2])
|
||||
assert not s.has_puzzle_subscription(phs[3])
|
||||
full_node_api.full_node.config["trusted_max_subscribe_items"] = 4
|
||||
full_node_api.full_node.config["trusted_peers"] = {server_2.node_id.hex(): server_2.node_id.hex()}
|
||||
assert full_node_api.is_trusted(con) is True
|
||||
msg_response = await full_node_api.register_interest_in_puzzle_hash(msg, con)
|
||||
assert msg_response.type == ProtocolMessageTypes.respond_to_ph_update.value
|
||||
assert len(s._ph_subscriptions) == 4
|
||||
assert s.has_ph_subscription(phs[0])
|
||||
assert s.has_ph_subscription(phs[1])
|
||||
assert s.has_ph_subscription(phs[2])
|
||||
assert s.has_ph_subscription(phs[3])
|
||||
assert not s.has_ph_subscription(phs[4])
|
||||
assert not s.has_ph_subscription(phs[5])
|
||||
assert s.puzzle_subscription_count() == 4
|
||||
assert s.has_puzzle_subscription(phs[0])
|
||||
assert s.has_puzzle_subscription(phs[1])
|
||||
assert s.has_puzzle_subscription(phs[2])
|
||||
assert s.has_puzzle_subscription(phs[3])
|
||||
assert not s.has_puzzle_subscription(phs[4])
|
||||
assert not s.has_puzzle_subscription(phs[5])
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_coin_subscribe_limits(self, simulator_and_wallet, self_hostname):
|
||||
@ -717,7 +717,7 @@ class TestSimpleSyncProtocol:
|
||||
msg_response = await full_node_api.register_interest_in_coin(msg, con)
|
||||
assert msg_response.type == ProtocolMessageTypes.respond_to_coin_update.value
|
||||
s = full_node_api.full_node.subscriptions
|
||||
assert len(s._coin_subscriptions) == 2
|
||||
assert s.coin_subscription_count() == 2
|
||||
assert s.has_coin_subscription(coins[0])
|
||||
assert s.has_coin_subscription(coins[1])
|
||||
assert not s.has_coin_subscription(coins[2])
|
||||
@ -727,7 +727,7 @@ class TestSimpleSyncProtocol:
|
||||
assert full_node_api.is_trusted(con) is True
|
||||
msg_response = await full_node_api.register_interest_in_coin(msg, con)
|
||||
assert msg_response.type == ProtocolMessageTypes.respond_to_coin_update.value
|
||||
assert len(s._coin_subscriptions) == 4
|
||||
assert s.coin_subscription_count() == 4
|
||||
assert s.has_coin_subscription(coins[0])
|
||||
assert s.has_coin_subscription(coins[1])
|
||||
assert s.has_coin_subscription(coins[2])
|
||||
|
||||
Loading…
Reference in New Issue
Block a user