Переглянути джерело

thread safe nonce creation

Lukas Goldschmidt 2 тижнів тому
батько
коміт
a7b0844998
3 змінених файлів з 68 додано та 0 видалено
  1. 0 0
      .codex
  2. 18 0
      src/exec_mcp/bitstamp.py
  3. 50 0
      tests/test_bitstamp_nonce.py

+ 18 - 0
src/exec_mcp/bitstamp.py

@@ -29,13 +29,31 @@ class AccountInfo:
 _AUTH_BREAKER_LOCK = threading.Lock()
 _AUTH_BREAKER_NEXT_ALLOWED: dict[str, float] = {}
 _AUTH_BREAKER_SECONDS = 5.0
+_NONCE_LOCK = threading.Lock()
+_LAST_NONCE_BY_SCOPE: dict[str, int] = {}
+
+
+def _nonce_now_ms() -> int:
+    return time.time_ns() // 1_000_000
+
+
+def _next_nonce(scope: str) -> int:
+    with _NONCE_LOCK:
+        last_nonce = _LAST_NONCE_BY_SCOPE.get(scope, 0)
+        nonce = max(_nonce_now_ms(), last_nonce + 1)
+        _LAST_NONCE_BY_SCOPE[scope] = nonce
+        return nonce
 
 
 class LG_Trading(Trading):
     def __init__(self, username, key, secret, *args, **kwargs):
         self._username = str(username)
+        self._nonce_scope = f"{username}:{key}"
         super(LG_Trading, self).__init__(username=username, key=key, secret=secret, *args, **kwargs)
 
+    def get_nonce(self):
+        return _next_nonce(self._nonce_scope)
+
     def _breaker_is_open(self) -> bool:
         with _AUTH_BREAKER_LOCK:
             return _AUTH_BREAKER_NEXT_ALLOWED.get(self._username, 0.0) > time.monotonic()

+ 50 - 0
tests/test_bitstamp_nonce.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from exec_mcp import bitstamp
+
+
+def test_next_nonce_advances_within_same_millisecond(monkeypatch):
+    bitstamp._LAST_NONCE_BY_SCOPE.clear()
+    monkeypatch.setattr(bitstamp, "_nonce_now_ms", lambda: 1_700_000_000_000)
+
+    first = bitstamp._next_nonce("acct-a:key-a")
+    second = bitstamp._next_nonce("acct-a:key-a")
+    third = bitstamp._next_nonce("acct-a:key-a")
+
+    assert first == 1_700_000_000_000
+    assert second == 1_700_000_000_001
+    assert third == 1_700_000_000_002
+
+
+def test_next_nonce_is_scoped_per_credentials(monkeypatch):
+    bitstamp._LAST_NONCE_BY_SCOPE.clear()
+    monkeypatch.setattr(bitstamp, "_nonce_now_ms", lambda: 1_700_000_000_000)
+
+    first_a = bitstamp._next_nonce("acct-a:key-a")
+    first_b = bitstamp._next_nonce("acct-b:key-b")
+    second_a = bitstamp._next_nonce("acct-a:key-a")
+
+    assert first_a == 1_700_000_000_000
+    assert first_b == 1_700_000_000_000
+    assert second_a == 1_700_000_000_001
+
+
+def test_lg_trading_instances_share_nonce_sequence(monkeypatch):
+    bitstamp._LAST_NONCE_BY_SCOPE.clear()
+    now_ms = [1_700_000_000_000]
+    monkeypatch.setattr(bitstamp, "_nonce_now_ms", lambda: now_ms[0])
+
+    first = bitstamp.LG_Trading.__new__(bitstamp.LG_Trading)
+    first._nonce_scope = "acct-a:key-a"
+
+    second = bitstamp.LG_Trading.__new__(bitstamp.LG_Trading)
+    second._nonce_scope = "acct-a:key-a"
+
+    nonce_1 = first.get_nonce()
+    nonce_2 = second.get_nonce()
+    now_ms[0] += 5
+    nonce_3 = first.get_nonce()
+
+    assert nonce_1 == 1_700_000_000_000
+    assert nonce_2 == 1_700_000_000_001
+    assert nonce_3 == 1_700_000_000_005