diff --git a/hathor/simulator/miner/geometric_miner.py b/hathor/simulator/miner/geometric_miner.py index 5887b2d94..53a0b8c0d 100644 --- a/hathor/simulator/miner/geometric_miner.py +++ b/hathor/simulator/miner/geometric_miner.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math from typing import TYPE_CHECKING, Optional from hathor.conf import HathorSettings @@ -50,6 +51,7 @@ def __init__( self._signal_bits = signal_bits or [] self._block: Optional[Block] = None self._blocks_found: int = 0 + self._blocks_before_pause: float = math.inf def _on_new_tx(self, key: HathorEvents, args: 'EventArguments') -> None: """ Called when a new tx or block is received. It updates the current mining to the @@ -81,12 +83,17 @@ def _generate_mining_block(self) -> 'Block': return block def _schedule_next_block(self): + if self._blocks_before_pause <= 0: + self._delayed_call = None + return + if self._block: self._block.nonce = self._rng.getrandbits(32) self._block.update_hash() self.log.debug('randomized step: found new block', hash=self._block.hash_hex, nonce=self._block.nonce) self._manager.propagate_tx(self._block, fails_silently=False) self._blocks_found += 1 + self._blocks_before_pause -= 1 self._block = None if self._manager.can_start_mining(): @@ -110,3 +117,16 @@ def _schedule_next_block(self): def get_blocks_found(self) -> int: return self._blocks_found + + def pause_after_exactly(self, *, n_blocks: int) -> None: + """ + Configure the miner to pause mining blocks after exactly `n_blocks` are propagated. If called more than once, + will unpause the miner and pause again according to the new argument. + + Use this instead of the `StopAfterNMinedBlocks` trigger if you need "exactly N blocks" behavior, instead of + "at least N blocks". + """ + self._blocks_before_pause = n_blocks + + if not self._delayed_call: + self._delayed_call = self._clock.callLater(0, self._schedule_next_block) diff --git a/hathor/simulator/trigger.py b/hathor/simulator/trigger.py index 2d54831f5..5745523ce 100644 --- a/hathor/simulator/trigger.py +++ b/hathor/simulator/trigger.py @@ -31,7 +31,12 @@ def should_stop(self) -> bool: class StopAfterNMinedBlocks(Trigger): - """Stop the simulation after `miner` finds N blocks. Note that these blocks might be orphan.""" + """ + Stop the simulation after `miner` finds at least N blocks. Note that these blocks might be orphan. + + Use `miner.pause_after_exactly()` instead of this trigger if you need "exactly N blocks" behavior, instead of + "at least N blocks". + """ def __init__(self, miner: 'AbstractMiner', *, quantity: int) -> None: self.miner = miner self.quantity = quantity diff --git a/tests/feature_activation/test_mining_simulation.py b/tests/feature_activation/test_mining_simulation.py index 1f1ec1354..d8e785a49 100644 --- a/tests/feature_activation/test_mining_simulation.py +++ b/tests/feature_activation/test_mining_simulation.py @@ -26,7 +26,6 @@ from hathor.feature_activation.settings import Settings as FeatureSettings from hathor.mining.ws import MiningWebsocketFactory, MiningWebsocketProtocol from hathor.p2p.resources import MiningResource -from hathor.simulator.trigger import StopAfterNMinedBlocks from hathor.transaction.resources import GetBlockTemplateResource from hathor.transaction.util import unpack, unpack_len from hathor.util import json_loadb @@ -88,47 +87,56 @@ def test_signal_bits_in_mining(self) -> None: # At the beginning, all features are outside their signaling period, so none are signaled. expected_signal_bits = 0b0000 assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1)) + miner.pause_after_exactly(n_blocks=1) + self.simulator.run(3600) assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=6)) + miner.pause_after_exactly(n_blocks=6) + self.simulator.run(3600) assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 6 # At height=8, NOP_FEATURE_1 is signaling, so it's enabled by the default support. expected_signal_bits = 0b0001 - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1)) + miner.pause_after_exactly(n_blocks=1) + self.simulator.run(3600) assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=3)) + miner.pause_after_exactly(n_blocks=3) + self.simulator.run(3600) assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 3 # At height=12, NOP_FEATURE_2 is signaling, enabled by the user. NOP_FEATURE_1 also continues signaling. expected_signal_bits = 0b0101 - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1)) + miner.pause_after_exactly(n_blocks=1) + self.simulator.run(3600) assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=7)) + miner.pause_after_exactly(n_blocks=7) + self.simulator.run(3600) assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 7 # At height=20, NOP_FEATURE_1 stops signaling, and NOP_FEATURE_2 continues. expected_signal_bits = 0b0100 - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1)) + miner.pause_after_exactly(n_blocks=1) + self.simulator.run(3600) assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=3)) + miner.pause_after_exactly(n_blocks=3) + self.simulator.run(3600) assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 3 # At height=24, all features have left their signaling period and therefore none are signaled. expected_signal_bits = 0b0000 - self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1)) + miner.pause_after_exactly(n_blocks=1) + self.simulator.run(3600) assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]