From 76ee0e2179b472c0ff9dedf81d0866f8d2a2eae9 Mon Sep 17 00:00:00 2001
From: shum <github.shum@liber.li>
Date: Thu, 22 Oct 2026 13:29:03 +0000
Subject: [PATCH] store: add balance turns

---
 .../src/StoreService/Names/Turns.hs           | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/apps/simplex-store-service/src/StoreService/Names/Turns.hs b/apps/simplex-store-service/src/StoreService/Names/Turns.hs
index b9a5560b6e..84a3a42e6c 100644
--- a/apps/simplex-store-service/src/StoreService/Names/Turns.hs
+++ b/apps/simplex-store-service/src/StoreService/Names/Turns.hs
@@ -1,22 +1,25 @@
 {-# LANGUAGE NamedFieldPuns #-}
 
--- | A thread that takes a turn it already holds waits forever. Label turns are taken before a code turn.
-module StoreService.Names.Turns (Turns, newTurns, withLabelTurn, withLabelTurns, withCodeTurn, turnsCount, labelTurnCount) where
+-- | A thread that takes a turn it already holds waits forever. Label turns are taken before a code turn or a balance turn,
+-- and a balance turn is never taken with a code turn.
+module StoreService.Names.Turns (Turns, newTurns, withLabelTurn, withLabelTurns, withCodeTurn, withBalanceTurn, turnsCount, labelTurnCount, balanceTurnCount) where
 
 import Control.Concurrent.STM
 import Control.Exception (bracket, bracket_)
 import Data.ByteString (ByteString)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
+import Data.Text.Encoding (encodeUtf8)
 import Simplex.Messaging.SimplexName (LabelHash (..))
+import StoreService.Names.Phases (RegistrationId (..))
 
 -- | Each lock counts the requests holding or waiting for it, and its entry goes when none are left.
 type TurnLocks = TVar (Map ByteString (TMVar (), Int))
 
-data Turns = Turns {labelTurns :: TurnLocks, codeTurns :: TurnLocks}
+data Turns = Turns {labelTurns :: TurnLocks, codeTurns :: TurnLocks, balanceTurns :: TurnLocks}
 
 newTurns :: IO Turns
-newTurns = Turns <$> newTVarIO Map.empty <*> newTVarIO Map.empty
+newTurns = Turns <$> newTVarIO Map.empty <*> newTVarIO Map.empty <*> newTVarIO Map.empty
 
 withLabelTurn :: Turns -> LabelHash -> IO a -> IO a
 withLabelTurn Turns {labelTurns} (LabelHash h) = withTurn labelTurns h
@@ -28,9 +31,9 @@ withLabelTurns turns a@(LabelHash ha) b@(LabelHash hb) action
   | ha < hb = withLabelTurn turns a $ withLabelTurn turns b action
   | otherwise = withLabelTurn turns b $ withLabelTurn turns a action
 
--- | The label and code hashes with a holder or a waiter.
+-- | The label hashes, code hashes and balances with a holder or a waiter.
 turnsCount :: Turns -> IO Int
-turnsCount Turns {labelTurns, codeTurns} = atomically $ (+) <$> (Map.size <$> readTVar labelTurns) <*> (Map.size <$> readTVar codeTurns)
+turnsCount Turns {labelTurns, codeTurns, balanceTurns} = atomically $ sum <$> mapM (fmap Map.size . readTVar) [labelTurns, codeTurns, balanceTurns]
 
 -- | The requests holding or waiting for the label hash's turn.
 labelTurnCount :: Turns -> LabelHash -> STM Int
@@ -40,6 +43,17 @@ labelTurnCount Turns {labelTurns} (LabelHash h) = maybe 0 snd . Map.lookup h <$>
 withCodeTurn :: Turns -> ByteString -> IO a -> IO a
 withCodeTurn Turns {codeTurns} = withTurn codeTurns
 
+-- | Keyed by the order the balance belongs to.
+withBalanceTurn :: Turns -> RegistrationId -> IO a -> IO a
+withBalanceTurn Turns {balanceTurns} = withTurn balanceTurns . balanceKey
+
+-- | The requests holding or waiting for the balance's turn.
+balanceTurnCount :: Turns -> RegistrationId -> STM Int
+balanceTurnCount Turns {balanceTurns} rid = maybe 0 snd . Map.lookup (balanceKey rid) <$> readTVar balanceTurns
+
+balanceKey :: RegistrationId -> ByteString
+balanceKey (RegistrationId rid) = encodeUtf8 rid
+
 withTurn :: TurnLocks -> ByteString -> IO a -> IO a
 withTurn locks key action =
   bracket (atomically enter) (const $ atomically leave) $ \lock ->