Skip to content

Add support for LeashID #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
axman6 merged 4 commits into from
Feb 19, 2026
Merged

Add support for LeashID #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c310018
Changing leashing to use optional id
axman6 Feb 5, 2026
78f1083
Add LeashID for client identification
axman6 Feb 5, 2026
04fa22c
Update tests and benchmarks
axman6 Feb 5, 2026
0f03998
Add optional LeashID to MsgRelease
axman6 Feb 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion ouroboros-network-protocols/bench-cddl/Main.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,7 +318,12 @@ localStateQueryMessages =
(Stateful.AnyMessage
StateIdle
(LocalStateQuery.MsgAcquire
(SpecificPoint (BlockPoint largeCBORBS))))
(SpecificPoint (BlockPoint largeCBORBS)) Nothing))
, AnyMessageWithResult
(Stateful.AnyMessage
StateIdle
(LocalStateQuery.MsgAcquire
(SpecificPoint (BlockPoint largeCBORBS)) (Just 2600)))
, AnyMessageWithResult
(Stateful.AnyMessage
StateAcquiring
Expand Down
8 changes: 4 additions & 4 deletions ouroboros-network-protocols/cddl/specs/local-state-query.cddl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,14 @@ failure = acquireFailurePointTooOld
query = any
result = any

msgAcquire = [0, base.point, ? bool]
/ [8, ? bool]
/ [10, ? bool]
msgAcquire = [0, point, ? word32]

@adithyaov adithyaov Feb 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are choosing to modify msgAcquire as opposed to adding a new message. Is this what we want to do?

@axman6 axman6 Feb 10, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the time being, yeah, this change is backwards compatible, but adding the new error type isn't. Either way, this probably isn't the implementation we'll use going forwards anyway.

/ [8, ? word32]
/ [10, ? word32]
msgAcquired = [1]
msgFailure = [2, failure]
msgQuery = [3, query]
msgResult = [4, result]
msgRelease = [5]
msgRelease = [5, ? word32]
msgReAcquire = [6, point]
/ [9]
/ [11]
Expand Down
13 changes: 7 additions & 6 deletions ouroboros-network-protocols/src/Ouroboros/Network/Protocol/LocalStateQuery/Client.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ newtype LocalStateQueryClient block point (query :: Type -> Type) m a =
--
data ClientStIdle block point query (m :: Type -> Type) a where
SendMsgAcquire :: Target point
-> Bool
-> Maybe LeashID
-> ClientStAcquiring block point query m a
-> ClientStIdle block point query m a

Expand Down Expand Up @@ -77,7 +77,8 @@ data ClientStAcquired block point query m a where
-> ClientStAcquiring block point query m a
-> ClientStAcquired block point query m a

SendMsgRelease :: m (ClientStIdle block point query m a)
SendMsgRelease :: Maybe LeashID
-> m (ClientStIdle block point query m a)
-> ClientStAcquired block point query m a

-- | In the 'StQuerying' protocol state, the client does not have agency.
Expand Down Expand Up @@ -127,8 +128,8 @@ mapLocalStateQueryClient fpoint fquery fresult =
-> ClientStAcquired block' point' query' m a
goAcquired (SendMsgQuery q k) = case fquery q of
Some q' -> SendMsgQuery q' (goQuerying q q' k)
goAcquired (SendMsgReAcquire tgt k) = SendMsgReAcquire (fpoint <$> tgt) (goAcquiring k)
goAcquired (SendMsgRelease k) = SendMsgRelease (fmap goIdle k)
goAcquired (SendMsgReAcquire tgt k) = SendMsgReAcquire (fpoint <$> tgt) (goAcquiring k)
goAcquired (SendMsgRelease mLeashId k) = SendMsgRelease mLeashId (fmap goIdle k)

goQuerying :: forall result result'.
query result
Expand Down Expand Up @@ -190,9 +191,9 @@ localStateQueryClientPeer (LocalStateQueryClient handler) =
Yield StateAcquired StateAcquiring
(MsgReAcquire tgt)
(handleStAcquiring stAcquiring)
SendMsgRelease stIdle ->
SendMsgRelease mLeashId stIdle ->
Yield StateAcquired StateIdle
MsgRelease
(MsgRelease mLeashId)
(Effect (handleStIdle <$> stIdle))

handleStQuerying
Expand Down
49 changes: 29 additions & 20 deletions ouroboros-network-protocols/src/Ouroboros/Network/Protocol/LocalStateQuery/Codec.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,27 +77,27 @@ codecLocalStateQuery version
State st
-> Message (LocalStateQuery block point query) st st'
-> CBOR.Encoding
encode _ (MsgAcquire (SpecificPoint pt) False) =
encode _ (MsgAcquire (SpecificPoint pt) Nothing) =
CBOR.encodeListLen 2
<> CBOR.encodeWord 0
<> encodePoint pt

encode _ (MsgAcquire (SpecificPoint pt) True) =
encode _ (MsgAcquire (SpecificPoint pt) (Just (LeashID leashId))) =
CBOR.encodeListLen 3
<> CBOR.encodeWord 0
<> encodePoint pt
<> CBOR.encodeBool True
<> CBOR.encodeWord32 leashId

encode _ (MsgAcquire VolatileTip False) =
encode _ (MsgAcquire VolatileTip Nothing) =
CBOR.encodeListLen 1
<> CBOR.encodeWord 8

encode _ (MsgAcquire VolatileTip True) =
encode _ (MsgAcquire VolatileTip (Just (LeashID leashId))) =
CBOR.encodeListLen 2
<> CBOR.encodeWord 8
<> CBOR.encodeBool True
<> CBOR.encodeWord32 leashId

encode _ (MsgAcquire ImmutableTip False)
encode _ (MsgAcquire ImmutableTip Nothing)
| canAcquireImmutable =
CBOR.encodeListLen 1
<> CBOR.encodeWord 10
Expand All @@ -106,11 +106,11 @@ codecLocalStateQuery version
++ "must be conditional on negotiating v16 of the node-to-client "
++ "protocol"

encode _ (MsgAcquire ImmutableTip True)
encode _ (MsgAcquire ImmutableTip (Just (LeashID leashId)))
| canAcquireImmutable =
CBOR.encodeListLen 2
<> CBOR.encodeWord 10
<> CBOR.encodeBool True
<> CBOR.encodeWord32 leashId
| otherwise =
error $ "encodeFailure: local state query: acquiring the immutable tip "
++ "must be conditional on negotiating v16 of the node-to-client "
Expand All @@ -135,10 +135,15 @@ codecLocalStateQuery version
<> CBOR.encodeWord 4
<> encodeResult query result

encode _ MsgRelease =
encode _ (MsgRelease Nothing) =
CBOR.encodeListLen 1
<> CBOR.encodeWord 5

encode _ (MsgRelease (Just (LeashID leashId))) =
CBOR.encodeListLen 2
<> CBOR.encodeWord 5
<> CBOR.encodeWord32 leashId

encode _ (MsgReAcquire (SpecificPoint pt)) =
CBOR.encodeListLen 2
<> CBOR.encodeWord 6
Expand Down Expand Up @@ -172,26 +177,26 @@ codecLocalStateQuery version
case (stok, f, len, key) of
(SingIdle, _, 2, 0) -> do
pt <- decodePoint
return (SomeMessage (MsgAcquire (SpecificPoint pt) False))
return (SomeMessage (MsgAcquire (SpecificPoint pt) Nothing))

(SingIdle, _, 3, 0) -> do
pt <- decodePoint
leashed <- CBOR.decodeBool
return (SomeMessage (MsgAcquire (SpecificPoint pt) leashed))
leashed <- CBOR.decodeWord32
return (SomeMessage (MsgAcquire (SpecificPoint pt) (Just (LeashID leashed))))

(SingIdle, _, 1, 8) -> do
return (SomeMessage (MsgAcquire VolatileTip False))
return (SomeMessage (MsgAcquire VolatileTip Nothing))

(SingIdle, _, 2, 8) -> do
leashed <- CBOR.decodeBool
return (SomeMessage (MsgAcquire VolatileTip leashed))
leashed <- CBOR.decodeWord32
return (SomeMessage (MsgAcquire VolatileTip (Just (LeashID leashed))))

(SingIdle, _, 1, 10) -> do
return (SomeMessage (MsgAcquire ImmutableTip False))
return (SomeMessage (MsgAcquire ImmutableTip Nothing))

(SingIdle, _, 2, 10) -> do
leashed <- CBOR.decodeBool
return (SomeMessage (MsgAcquire ImmutableTip leashed))
leashed <- CBOR.decodeWord32
return (SomeMessage (MsgAcquire ImmutableTip (Just (LeashID leashed))))

(SingAcquiring, _, 1, 1) ->
return (SomeMessage MsgAcquired)
Expand All @@ -209,7 +214,11 @@ codecLocalStateQuery version
return (SomeMessage (MsgResult result))

(SingAcquired, _, 1, 5) ->
return (SomeMessage MsgRelease)
return (SomeMessage $ MsgRelease Nothing)

(SingAcquired, _, 2, 5) -> do
leashId <- CBOR.decodeWord32
return (SomeMessage $ MsgRelease (Just (LeashID leashId)))

(SingAcquired, _, 2, 6) -> do
pt <- decodePoint
Expand Down
23 changes: 12 additions & 11 deletions ouroboros-network-protocols/src/Ouroboros/Network/Protocol/LocalStateQuery/Server.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ newtype LocalStateQueryServer block point (query :: Type -> Type) m a = LocalSta
--
data ServerStIdle block point query m a = ServerStIdle {
recvMsgAcquire :: Target point
-> Bool
-> Maybe LeashID
-> m (ServerStAcquiring block point query m a),

recvMsgDone :: m a
Expand Down Expand Up @@ -74,7 +74,8 @@ data ServerStAcquired block point query m a = ServerStAcquired {
recvMsgReAcquire :: Target point
-> m (ServerStAcquiring block point query m a),

recvMsgRelease :: m (ServerStIdle block point query m a)
recvMsgRelease :: Maybe LeashID
-> m (ServerStIdle block point query m a)
}

-- | In the 'StQuerying' protocol state, the server has agency and must send:
Expand Down Expand Up @@ -127,15 +128,15 @@ localStateQueryServerPeer (LocalStateQueryServer handler) =
-> Server (LocalStateQuery block point query) StAcquired State m a
handleStAcquired ServerStAcquired{recvMsgQuery, recvMsgReAcquire, recvMsgRelease} =
Await $ \_ req -> case req of
MsgQuery query -> ( Effect $ handleStQuerying query <$> recvMsgQuery query
, StateQuerying query
)
MsgReAcquire pt -> ( Effect $ handleStAcquiring <$> recvMsgReAcquire pt
, StateAcquiring
)
MsgRelease -> ( Effect $ handleStIdle <$> recvMsgRelease
, StateIdle
)
MsgQuery query -> ( Effect $ handleStQuerying query <$> recvMsgQuery query
, StateQuerying query
)
MsgReAcquire pt -> ( Effect $ handleStAcquiring <$> recvMsgReAcquire pt
, StateAcquiring
)
MsgRelease mLeashId -> ( Effect $ handleStIdle <$> recvMsgRelease mLeashId
, StateIdle
)

handleStQuerying
:: query result
Expand Down
28 changes: 22 additions & 6 deletions ouroboros-network-protocols/src/Ouroboros/Network/Protocol/LocalStateQuery/Type.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- | The type of the local ledger state query protocol.
--
Expand All @@ -29,6 +31,8 @@ import Network.TypedProtocol.Stateful.Codec (AnyMessage (..))
import Control.DeepSeq
import GHC.Generics
import Ouroboros.Network.Util.ShowProxy (ShowProxy (..))
import Data.Word (Word32)
import NoThunks.Class (NoThunks)


-- | The kind of the local state query protocol, and the types of
Expand Down Expand Up @@ -128,6 +132,12 @@ data Target point = -- | The tip of the volatile chain
| ImmutableTip
deriving (Eq, Foldable, Functor, Generic, Ord, Show, Traversable, NFData)


newtype LeashID = LeashID Word32
deriving stock (Show)
-- TODO: anything else?
deriving newtype (Eq, Ord, NFData, Num, Read, NoThunks)

instance Protocol (LocalStateQuery (block :: Type) (point :: Type) (query :: Type -> Type)) where

-- | The messages in the state query protocol.
Expand All @@ -141,7 +151,7 @@ instance Protocol (LocalStateQuery (block :: Type) (point :: Type) (query :: Typ
--
MsgAcquire
:: Target point
-> Bool
-> Maybe LeashID -- ^ Optional leashing ID
-> Message (LocalStateQuery block point query) StIdle StAcquiring

-- | The server can confirm that it has the state at the requested point.
Expand Down Expand Up @@ -169,10 +179,14 @@ instance Protocol (LocalStateQuery (block :: Type) (point :: Type) (query :: Typ
-> Message (LocalStateQuery block point query) (StQuerying result) StAcquired

-- | The client can instruct the server to release the state. This lets
-- the server free resources.
-- the server free resources. If the LeashID is set, it means that this client's
-- leash should be removed, otherwise the node should continue to be leashed
-- for later reconnection.
--
MsgRelease
:: Message (LocalStateQuery block point query) StAcquired StIdle
-- TODO: Should there be a backwards compatible pattern and a new one for leashing?
:: Maybe LeashID -- If this is set, then it means the client wishes to UNleash
-> Message (LocalStateQuery block point query) StAcquired StIdle

-- | This is like 'MsgAcquire' but for when the client already has a
-- state. By moving to another state directly without a 'MsgRelease' it
Expand Down Expand Up @@ -210,7 +224,7 @@ instance ( forall result. NFData (query result)
rnf (MsgFailure af) = rnf af
rnf (MsgQuery qr) = rnf qr
rnf (MsgResult r) = rwhnf r
rnf MsgRelease = ()
rnf (MsgRelease lId) = rnf lId
rnf (MsgReAcquire mbPoint) = rnf mbPoint
rnf MsgDone = ()

Expand Down Expand Up @@ -254,8 +268,10 @@ instance (ShowQuery query, Show point)
showParen (p >= 11) $
showString "MsgResult " .
showParen True (showString (showResult query result))
AnyMessage _f MsgRelease ->
showString "MsgRelease"
AnyMessage _f (MsgRelease mLeashId) ->
showParen (p >= 11) $
showString "MsgRelease" .
showsPrec 11 mLeashId
AnyMessage _f (MsgReAcquire pt) ->
showParen (p >= 11) $
showString "MsgReAcquire " .
Expand Down
4 changes: 2 additions & 2 deletions ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/LocalStateQuery/Direct.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,9 @@ direct (LocalStateQueryClient mclient) (LocalStateQueryServer mserver) = do
directAcquired (SendMsgReAcquire tgt client') ServerStAcquired{recvMsgReAcquire} = do
server' <- recvMsgReAcquire tgt
directAcquiring client' server'
directAcquired (SendMsgRelease client) ServerStAcquired{recvMsgRelease} = do
directAcquired (SendMsgRelease mLeashId client) ServerStAcquired{recvMsgRelease} = do
client' <- client
server' <- recvMsgRelease
server' <- recvMsgRelease mLeashId
directIdle client' server'

directQuerying
Expand Down
Loading
Loading