-- | Connection lifecycle for the Switch.
--
-- Provides the single teardown path for upgraded connections and a
-- resource-accounted outbound stream opener. Teardown fires from the
-- stream accept loop exit (remote disconnect or session death), from
-- explicit 'closeConnection' calls, and from 'switchClose' via
-- 'closeAllConnections'.
module LibP2P.Switch.Connection
  ( closeConnection
  , closeAllConnections
  , newStream
  ) where

import Control.Concurrent.STM (atomically, readTVar, writeTChan, writeTVar)
import Control.Exception (SomeException, catch, finally, onException)
import Control.Monad (unless, when)
import Data.IORef (atomicModifyIORef', newIORef)
import LibP2P.MultistreamSelect.Negotiation (StreamIO (..))
import LibP2P.Switch.ConnPool (allConns, removeConn)
import LibP2P.Switch.ResourceManager
  ( Direction (..)
  , ResourceError
  , releaseConnection
  , releasePeerStream
  , reservePeerStream
  )
import LibP2P.Switch.Types
  ( ConnState (..)
  , Connection (..)
  , MuxerSession (..)
  , Switch (..)
  , SwitchEvent (..)
  )

-- | Tear down a connection: remove it from the pool, release its
-- resource reservation, publish a Disconnected event, and close the
-- muxer session together with the underlying transport.
--
-- Idempotent: the state transition to ConnClosed is atomic, so
-- concurrent calls (accept loop exit, explicit close, switchClose)
-- perform the teardown exactly once.
closeConnection :: Switch -> Connection -> IO ()
closeConnection :: Switch -> Connection -> IO ()
closeConnection Switch
sw Connection
conn = do
  shouldClose <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
    st <- TVar ConnState -> STM ConnState
forall a. TVar a -> STM a
readTVar (Connection -> TVar ConnState
connState Connection
conn)
    if st == ConnClosed
      then pure False
      else do
        writeTVar (connState conn) ConnClosed
        removeConn (swConnPool sw) conn
        releaseConnection (swResourceMgr sw) (connPeerId conn) (connDirection conn)
        writeTChan (swEvents sw)
          (Disconnected (connPeerId conn) (connDirection conn) (connRemoteAddr conn))
        pure True
  when shouldClose $
    muxClose (connSession conn) `catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Tear down every pooled connection (used by switchClose).
closeAllConnections :: Switch -> IO ()
closeAllConnections :: Switch -> IO ()
closeAllConnections Switch
sw = do
  conns <- STM [Connection] -> IO [Connection]
forall a. STM a -> IO a
atomically (STM [Connection] -> IO [Connection])
-> STM [Connection] -> IO [Connection]
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId [Connection]) -> STM [Connection]
allConns (Switch -> TVar (Map PeerId [Connection])
swConnPool Switch
sw)
  mapM_ (closeConnection sw) conns

-- | Open an outbound stream on a connection, reserving a stream slot
-- against the peer's resource scope. The slot is released when the
-- returned stream is closed (exactly once, even on double close).
newStream :: Switch -> Connection -> IO (Either ResourceError StreamIO)
newStream :: Switch -> Connection -> IO (Either ResourceError StreamIO)
newStream Switch
sw Connection
conn = do
  let pid :: PeerId
pid = Connection -> PeerId
connPeerId Connection
conn
      release :: IO ()
release = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ ResourceManager -> PeerId -> Direction -> STM ()
releasePeerStream (Switch -> ResourceManager
swResourceMgr Switch
sw) PeerId
pid Direction
Outbound
  reserved <- STM (Either ResourceError ()) -> IO (Either ResourceError ())
forall a. STM a -> IO a
atomically (STM (Either ResourceError ()) -> IO (Either ResourceError ()))
-> STM (Either ResourceError ()) -> IO (Either ResourceError ())
forall a b. (a -> b) -> a -> b
$ ResourceManager
-> PeerId -> Direction -> STM (Either ResourceError ())
reservePeerStream (Switch -> ResourceManager
swResourceMgr Switch
sw) PeerId
pid Direction
Outbound
  case reserved of
    Left ResourceError
err -> Either ResourceError StreamIO -> IO (Either ResourceError StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResourceError -> Either ResourceError StreamIO
forall a b. a -> Either a b
Left ResourceError
err)
    Right () -> do
      stream <- MuxerSession -> IO StreamIO
muxOpenStream (Connection -> MuxerSession
connSession Connection
conn) IO StreamIO -> IO () -> IO StreamIO
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release
      releasedRef <- newIORef False
      let releaseOnce = do
            already <- IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
releasedRef (\Bool
r -> (Bool
True, Bool
r))
            unless already release
      pure (Right stream { streamClose = streamClose stream `finally` releaseOnce })