{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Connection state, packet I/O, and the incremental token reader.
--
-- Everything here is protected by the connection's reentrant mutex.  A request
-- creates an 'ActiveResponse'; tokens are pulled from it one at a time, and
-- starting a new request first cancels and drains any unfinished response so
-- the connection can never be left mid-message.
--
-- A request and its response travel on a 'Session'.  A connection has one,
-- 'connPrimary', unless it negotiated MARS ('connMars'): then each request
-- takes an SMP session of its own ("Hmssql.Tds.Smp"), so several responses
-- can be open at once, and a new request leaves the others alone.  The
-- backend, the transport, the protocol state ('connContext', the transaction
-- descriptor included), the mutex and the broken flag stay the
-- connection's: the sessions are one SQL session on the server.
--
-- Interruptions (asynchronous exceptions) are handled by two rules:
--
-- * A send that doesn't complete breaks the connection ('sendBytes'): the
--   server may hold part of a message, and nothing here tries to recover
--   from that.
--
-- * A read that is interrupted leaves the connection usable ('nextToken'):
--   an Attention is sent at once, which cancels the statement, and the next
--   request drains the canceled response up to the server's
--   acknowledgement ('finishResponse').  If that drain is itself
--   interrupted, the connection is broken too.
--
-- A 'Connection' value is a handle, a lease on the connection: every field
-- but 'connLease' is shared or immutable, so a record copy is another handle
-- to the same connection.  'Hmssql.Connection.release' invalidates the
-- handle it was given ('invalidateLease') and 'renewLease' makes a new one;
-- everything that takes the lock, and every row a stream yields, checks the
-- handle first ('checkLease').
--
-- Applications use "Hmssql.Connection"; nothing here is stable.
module Hmssql.Connection.Internal
  ( Connection (..)
  , Transport (..)
  , Session (..)
  , SmpSession (..)
  , SmpLink (..)
  , Mars (..)
  , ActiveResponse (..), Cancellation (..)
  , ReadState (..), ReadPosition (..)
  , newConnection
  , startReader
  , enableMars
  , allSessions
  , closeConnection
  , releaseTransport
  , closeSocket
  , closeSocketWith
  , socketCloser
  , connectionContext
  , connectionMars
  , withConnectionLock
  , checkLease
  , invalidateLease
  , renewLease
  , renewUsable
  , checkNotBroken
  , refuseUnusable
  , currentUnusableReason
  , markBroken
  , reasonFor
  , resetFailure
  , throwSevere
  , transportError
  , recvSome
  , recvExactly
  , newBufferedRecv
  , SendGather (..)
  , Gathered (..)
  , gatheredSend
  , sendBytes
  , sendFramed
  , sendPackets
  , readPacket
  , sendMessage
  , runBody
  , startRequest
  , sourceAction
  , bytesBody
  , startRequestFor
  , startStreamingRequest
  , withReservedSession
  , requestReset
  , resetWithSessionSetup
  , RowRead (..), newRowRead, wholeRows, discardRows, drainingRows
  , nextStepReversed
  , RowBatch (..)
  , ResponseStep (..)
  , nextStep
  , readPiece
  , nextToken
  , checkCurrent
  , finishResponse
  , cancelActiveResponse
  , readingToEnd
  , abandonResponse
  , drainCancelled
  , busyStatements
  , drainActiveResponse
  , readWholeMessage
  , readWholeMessageUpTo
  ) where

import Control.Concurrent (ThreadId, forkIOWithUnmask, forkOnWithUnmask, killThread, mkWeakThreadId, myThreadId, threadCapability)
import GHC.Conc (labelThread)
import GHC.Exts (keepAlive#)
import GHC.IO (IO (..))
import Control.Concurrent.MVar
import Control.Concurrent.STM
import Control.Exception
import Control.Applicative ((<|>))
import Control.Monad (forM, forM_, unless, void, when)
import Streaming (Stream, inspect)
import qualified Streaming
import Hmssql.Internal.Derived
import Hmssql.Tds.Reader
import qualified Hmssql.Tds.Value as V
import qualified Hmssql.Tds.Parser as P
import Data.Binary.Put (Put)
import qualified Data.Binary.Put as Put
import Data.Bits ((.&.))
import Data.Int (Int32)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Builder.Extra as BB
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Unsafe as BU
import Data.IORef
import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing)
import Data.Proxy (Proxy (..))
import Data.Text (Text)
import qualified Data.Text as T
import Data.Word (Word16, Word32, Word8)
import Foreign.ForeignPtr (ForeignPtr, plusForeignPtr, withForeignPtr)
import Foreign.Marshal.Utils (copyBytes, moveBytes)
import Foreign.Ptr (plusPtr)
import qualified Network.TLS as TLS
import System.Mem.Weak (Weak, deRefWeak)

import Hmssql.Connection.Backend
import Hmssql.Connection.Locking
import Hmssql.Errors
import Hmssql.Tds.Client (sqlBatchBuilder, sqlBatchProblem)
import Hmssql.Tds.Context
import Hmssql.Tds.Packet
import Hmssql.Tds.PacketWriter
import Hmssql.Tds.Smp
import Hmssql.Tds.Token

-- | How the connection's bytes go to the backend: as they are, or through
-- a TLS session.
data Transport
  = TransportPlain !Backend
  | TransportTls !TLS.Context

-- | What one request at a time travels on: the connection itself, or under
-- MARS one of its SMP sessions.  A session is /busy/ while it has a current
-- response, that is, until the response's end has been read ('StepEnd'),
-- whether by its consumer or by a drain after it was canceled.
data Session = Session
  { Session -> IORef (Maybe ActiveResponse)
sessResponse :: !(IORef (Maybe ActiveResponse))
    -- ^ The response being read on the session, if any.
  , Session -> IORef Int
sessGeneration :: !(IORef Int)
    -- ^ Counts the session's responses; one whose 'arGeneration' differs
    -- from its session's current response's was superseded or canceled.
  , Session -> IORef [ByteString]
sessRecvBuffer :: !(IORef [B.ByteString])
    -- ^ Without SMP: bytes read from the transport but not yet consumed
    -- as a whole packet: non-empty chunks, oldest first.  Kept as they
    -- arrived, so only a packet split between chunks is copied
    -- ('takeBytes').
  , Session -> Maybe SmpSession
sessSmp :: !(Maybe SmpSession)
    -- ^ The SMP session, on a MARS connection once SMP has started.
  }

-- | An SMP session's state on the side of the threads that use the
-- connection.
data SmpSession = SmpSession
  { SmpSession -> Word16
smpSid :: !Word16
  , SmpSession -> SmpLink
smpLink :: !SmpLink
  , SmpSession -> IORef Word32
smpSendSeq :: !(IORef Word32)
    -- ^ The sequence number of the last DATA packet sent (SeqNumForSend).
    -- Written under the send lock.
  , SmpSession -> IORef Word32
smpAdvertised :: !(IORef Word32)
    -- ^ The window last sent to the server.
  }

-- | An SMP session's state shared with the reader thread, which refers to
-- nothing else of the session (see 'startReader').
data SmpLink = SmpLink
  { SmpLink -> TQueue ByteString
linkIncoming :: !(TQueue B.ByteString)
    -- ^ The TDS packets received, whole, in order.  At most a window of
    -- them: the server sends no further than it was told.
  , SmpLink -> TVar Word32
linkSendHighWater :: !(TVar Word32)
    -- ^ The highest sequence number the server accepts (HighWaterForSend):
    -- the window of the last packet received on the session.
  , SmpLink -> TVar Word32
linkRecvSeq :: !(TVar Word32)
    -- ^ The sequence number of the last DATA packet received
    -- (SeqNumForRecv).
  , SmpLink -> TVar Word32
linkRecvHighWater :: !(TVar Word32)
    -- ^ The highest sequence number accepted from the server
    -- (HighWaterForRecv): the packets consumed plus the window.  Raised
    -- by the consumer before it tells the server.
  , SmpLink -> TVar Int
linkEoms :: !(TVar Int)
    -- ^ How many packets marked end-of-message have been received.
  , SmpLink -> TVar Bool
linkClosing :: !(TVar Bool)
    -- ^ A FIN was sent; the server's FIN in reply frees the identifier.
  }

-- | What a MARS connection has beyond a plain one.
data Mars = Mars
  { Mars -> TVar (IntMap SmpLink)
marsRegistry :: !(TVar (IntMap.IntMap SmpLink))
    -- ^ The sessions the server knows, by identifier: open, or closing (a
    -- FIN sent and not yet answered).  Shared with the reader thread.
  , Mars -> IORef (IntMap Session)
marsOpen :: !(IORef (IntMap.IntMap Session))
    -- ^ The open sessions.
  , Mars -> IORef (Maybe Session)
marsReserved :: !(IORef (Maybe Session))
    -- ^ The session every request is sent on for now
    -- ('withReservedSession').
  , Mars -> Word32
marsWindow :: !Word32
    -- ^ How many packets the server may send ahead of what was consumed,
    -- on each session.
  , Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight :: !(IORef (Maybe ActiveResponse))
    -- ^ The response to a request that carried RESETCONNECTION, until its
    -- leading ENVCHANGEs have been read: no request starts before
    -- ('awaitReset').
  }

-- | The server's reply to one request.
data ActiveResponse = ActiveResponse
  { ActiveResponse -> Int
arGeneration :: !Int
  , ActiveResponse -> Session
arSession :: !Session
    -- ^ What the request was sent on.
  , ActiveResponse -> Maybe Text
arStatement :: !(Maybe Text)
    -- ^ The request's SQL text, when known: to name an unfinished response.
  , ActiveResponse -> Int
arEomBase :: !Int
    -- ^ The session's 'linkEoms' when the request was sent: the response
    -- has fully arrived once the count is past it.
  , ActiveResponse -> Bool
arCarriedReset :: !Bool
    -- ^ The request carried RESETCONNECTION.
  , ActiveResponse -> IORef Bool
arSawResetAck :: !(IORef Bool)
    -- ^ The server's acknowledgement of a reset was read in this response.
  , ActiveResponse -> IORef Bool
arInConsumer :: !(IORef Bool)
    -- ^ Reading this response recursively from its value consumer is refused.
  , ActiveResponse -> IORef Cancellation
arCancellation :: !(IORef Cancellation)
    -- ^ Cancellation and abandonment have one progression, rather than
    -- independent flags admitting contradictory states.
  , ActiveResponse -> IORef [Token]
arReplay :: !(IORef [Token])
    -- ^ Tokens read ahead of the consumer by 'awaitReset', in order, with
    -- their effects already applied: the consumer is given them first.
  , ActiveResponse -> IORef ReadState
arState :: !(IORef ReadState)
    -- ^ What has been received but not yet returned.
  , ActiveResponse -> IORef (Maybe ResultSetInfo)
arResultSet :: !(IORef (Maybe ResultSetInfo))
    -- ^ The columns of the most recent COLMETADATA, compiled for the row
    -- recipes and decoders.
  , ActiveResponse -> IORef Bool
arSawFinalDone :: !(IORef Bool)
  }

data Cancellation = Running | Abandoned | AttentionSent | AttentionAcknowledged
  deriving stock (Cancellation -> Cancellation -> Bool
(Cancellation -> Cancellation -> Bool)
-> (Cancellation -> Cancellation -> Bool) -> Eq Cancellation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Cancellation -> Cancellation -> Bool
== :: Cancellation -> Cancellation -> Bool
$c/= :: Cancellation -> Cancellation -> Bool
/= :: Cancellation -> Cancellation -> Bool
Eq, Int -> Cancellation -> ShowS
[Cancellation] -> ShowS
Cancellation -> [Char]
(Int -> Cancellation -> ShowS)
-> (Cancellation -> [Char])
-> ([Cancellation] -> ShowS)
-> Show Cancellation
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Cancellation -> ShowS
showsPrec :: Int -> Cancellation -> ShowS
$cshow :: Cancellation -> [Char]
show :: Cancellation -> [Char]
$cshowList :: [Cancellation] -> ShowS
showList :: [Cancellation] -> ShowS
Show)

attentionWasSent :: Cancellation -> Bool
attentionWasSent :: Cancellation -> Bool
attentionWasSent Cancellation
AttentionSent = Bool
True
attentionWasSent Cancellation
AttentionAcknowledged = Bool
True
attentionWasSent Cancellation
_ = Bool
False

markAbandoned :: ActiveResponse -> IO ()
markAbandoned :: ActiveResponse -> IO ()
markAbandoned ActiveResponse
ar = IORef Cancellation -> (Cancellation -> Cancellation) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (ActiveResponse -> IORef Cancellation
arCancellation ActiveResponse
ar) ((Cancellation -> Cancellation) -> IO ())
-> (Cancellation -> Cancellation) -> IO ()
forall a b. (a -> b) -> a -> b
$ \case
  Cancellation
Running -> Cancellation
Abandoned
  Cancellation
state -> Cancellation
state

-- | What has been received but not yet returned: whether the message's last
-- packet has arrived, and where the parse has got to.
data ReadState = ReadState
  { ReadState -> Bool
readFinalPacket :: !Bool
  , ReadState -> ReadPosition
readPosition :: !ReadPosition
  }

-- Where the response's parse has got to.  The rest of the parse is the state, so
-- an interrupted read, a later pull and a drain all go on from the same value,
-- and nothing has to be rewound or decoded twice; a column and a phase stored
-- here instead would need a parser of their own for the drain.
data ReadPosition
  = AtToken !B.ByteString
    -- ^ Between tokens, with the unconsumed payload.
  | Reading (Stream P.Await IO ReadEvent)
    -- ^ Normalized suspension, ready for a packet or a consumer outcome.
  | Parked !B.ByteString (P.Parser ())
    -- ^ A generic cell pull returned; the recipe drains its remaining row.

-- A completed suspended read either supplies a token, delivers to its typed
-- row recipient, or reports a skipped row. Every event includes unread bytes.
--
-- The response outlives any one decoder, and queryMulti changes decoder between
-- result sets, so a suspension stored here can't mention the row's type.  So a
-- suspended row ends in 'ReadRow', whose action puts the result in the slot of the
-- 'RowRead' that started it.  A row that completes without suspending never touches
-- the slot; it comes back from 'P.runParser' directly.
data ReadEvent
  = ReadToken !Token !B.ByteString
  | ReadRow (IO ()) !B.ByteString
  | SkippedRow !B.ByteString

-- | An open connection to a server.  It belongs to one thread at a time;
-- "Hmssql.Connection" has the rules.
data Connection = Connection
  { Connection -> Backend
connBackend :: !Backend
    -- ^ The byte stream the connection runs on: a socket, or on Windows a
    -- pipe ("Hmssql.Connection.Backend").
  , Connection -> IORef Transport
connTransport :: !(IORef Transport)
  , Connection -> Session
connPrimary :: !Session
    -- ^ The session of a connection without MARS, and of any connection
    -- until it has logged in.  Under MARS, the SMP session 0.
  , Connection -> Maybe Mars
connMars :: !(Maybe Mars)
    -- ^ Set once SMP has started ('enableMars').
  , Connection -> IORef ()
connAlive :: !(IORef ())
    -- ^ What the finalizer that stops the reader thread and closes the
    -- backend is attached to ('startReader'): every handle reaches it, and
    -- the reader thread doesn't.
  , Connection -> TMVar ByteString
connIncoming :: !(TMVar B.ByteString)
    -- ^ Chunks from the reader thread, once it is running (without MARS;
    -- under MARS it puts packets in their sessions' 'linkIncoming').
  , Connection -> TVar (Maybe SomeException)
connEnded :: !(TVar (Maybe SomeException))
    -- ^ Why the transport can no longer be used: set when the reader thread
    -- exits (to the 'ConnectionClosed' it saw, or 'ConnectionBroken' if it
    -- died of anything else), or by 'closeConnection'.  Reads see it once
    -- 'connIncoming' is drained, or at once after 'closeConnection'.
  , Connection -> MVar ()
connSendLock :: !(MVar ())
    -- ^ Held while sending, so 'closeConnection' can tell whether it may say
    -- goodbye to the server.
  , Connection -> SendGather
connSendGather :: !SendGather
    -- ^ The TLS records of the message being sent, gathered into socket
    -- writes ('gatheredSend').
  , Connection -> IORef (Maybe (Weak ThreadId))
connReader :: !(IORef (Maybe (Weak ThreadId)))
    -- ^ The reader thread, once it is running ('startReader').  Weak: a
    -- 'ThreadId' keeps the thread's stack alive (32 KiB once it has grown)
    -- even after the thread has finished, which a closed connection that is
    -- still referenced shouldn't.  A running reader is reachable anyway
    -- (it is blocked on the socket or on 'connIncoming').
  , Connection -> IORef TdsContext
connContext :: !(IORef TdsContext)
  , Connection -> Mutex
connMutex :: !Mutex
  , Connection -> IORef (Maybe UnusableReason)
connBroken :: !(IORef (Maybe UnusableReason))
  , Connection -> IORef Bool
connResetPending :: !(IORef Bool)
    -- ^ Set by 'requestReset'; consumed by the next SQLBatch or RPC.
  , Connection -> Bool
connPinReader :: !Bool
    -- ^ Whether 'startReader' pins the reader thread to the capability
    -- of the thread that connects ('Hmssql.Connection.connectPinReader').
  , Connection -> [Text]
connSessionSetup :: ![Text]
    -- ^ Run after login (as one batch), and again after each reset.
  , Connection -> Int
connValueAllocAhead :: !Int
    -- ^ How far ahead of its bytes a large value's buffer is allocated
    -- ('Hmssql.Connection.connectValueAllocAhead').
  , Connection -> ServerMessage -> IO ()
connOnInfo :: !(ServerMessage -> IO ())
  , Connection -> EnvChange -> IO ()
connOnEnvChange :: !(EnvChange -> IO ())
  , Connection -> IORef Bool
connInSource :: !(IORef Bool)
  , Connection -> IORef Bool
connInCallback :: !(IORef Bool)
    -- ^ A callback is running.  It runs with the lock held, in the middle of
    -- reading a response, so it must not use the connection.
  , Connection -> Int
connLease :: !Int
    -- ^ This handle's lease: the only field not shared between handles
    -- to the connection.
  , Connection -> IORef Int
connCurrentLease :: !(IORef Int)
    -- ^ The connection's current lease.  A handle whose 'connLease'
    -- differs was released, or superseded by 'renewLease' ('checkLease').
  }

-- | A connection with no session setup; set 'connSessionSetup' afterwards.
newConnection :: Backend -> TdsContext -> (ServerMessage -> IO ()) -> IO Connection
newConnection :: Backend -> TdsContext -> (ServerMessage -> IO ()) -> IO Connection
newConnection Backend
backend TdsContext
ctx ServerMessage -> IO ()
onInfo = do
  connTransport <- Transport -> IO (IORef Transport)
forall a. a -> IO (IORef a)
newIORef (Backend -> Transport
TransportPlain Backend
backend)
  connPrimary <- newSession Nothing
  connAlive <- newIORef ()
  connIncoming <- newEmptyTMVarIO
  connEnded <- newTVarIO Nothing
  connSendLock <- newMVar ()
  connSendGather <- newSendGather
  connReader <- newIORef Nothing
  connContext <- newIORef ctx
  connMutex <- newMutex
  connBroken <- newIORef Nothing
  connInCallback <- newIORef False
  connInSource <- newIORef False
  connResetPending <- newIORef False
  connCurrentLease <- newIORef 0
  pure Connection
    { connBackend = backend
    , connSessionSetup = []
    , connPinReader = False
    , connValueAllocAhead = lobAllocAhead
    , connOnInfo = onInfo
    , connOnEnvChange = \EnvChange
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    , connLease = 0
    , connMars = Nothing
    , ..
    }

newSession :: Maybe SmpSession -> IO Session
newSession :: Maybe SmpSession -> IO Session
newSession Maybe SmpSession
sessSmp = do
  sessResponse <- Maybe ActiveResponse -> IO (IORef (Maybe ActiveResponse))
forall a. a -> IO (IORef a)
newIORef Maybe ActiveResponse
forall a. Maybe a
Nothing
  sessGeneration <- newIORef 0
  sessRecvBuffer <- newIORef []
  pure Session{..}

-- | The protocol state negotiated for the connection: the TDS version, the
-- packet size the server granted, whether a transaction is open, the
-- session's collation.
connectionContext :: Connection -> IO TdsContext
connectionContext :: Connection -> IO TdsContext
connectionContext = IORef TdsContext -> IO TdsContext
forall a. IORef a -> IO a
readIORef (IORef TdsContext -> IO TdsContext)
-> (Connection -> IORef TdsContext) -> Connection -> IO TdsContext
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> IORef TdsContext
connContext

-- | Whether the connection negotiated MARS ('Hmssql.Connection.connectMars'):
-- several responses can be open on it at once, and a new request leaves an
-- open stream alone.  For code that is handed a 'Connection' without the
-- 'Hmssql.Connection.ConnectInfo' it was opened with.
connectionMars :: Connection -> Bool
connectionMars :: Connection -> Bool
connectionMars = Maybe Mars -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Mars -> Bool)
-> (Connection -> Maybe Mars) -> Connection -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> Maybe Mars
connMars

-- | Hold the connection's reentrant lock.  Refuses (with 'ReentrantCallback')
-- a callback's attempt to use the connection: the callback's thread already
-- holds the lock, so the attempt would otherwise interleave with the response
-- being read.  Refuses a released handle too ('checkLease'), once the lock
-- is held, so a handle released while this call waited for the lock is
-- refused.
withConnectionLock :: Connection -> IO a -> IO a
withConnectionLock :: forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn IO a
act = Mutex -> IO a -> IO a
forall a. Mutex -> IO a -> IO a
withMutex (Connection -> Mutex
connMutex Connection
conn) (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
  inCallback <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
connInCallback Connection
conn)
  when inCallback $ throwIO ReentrantCallback
  source <- readIORef (connInSource conn)
  when (source && isNothing (connMars conn)) (throwIO ReentrantSource)
  checkLease conn
  act
-- Opaque, so that GHC doesn't specialize the lock, the lease check and the
-- callback check into each caller's action: called, a request allocates 128
-- bytes less (100 executes against the mock server, 1,213,896 bytes in place
-- of 1,226,696, GHC 9.12.4, 2026-11-04).
{-# OPAQUE withConnectionLock #-}

-- | User code producing request data may not start a request on this connection.
sourceAction :: Connection -> IO a -> IO a
sourceAction :: forall a. Connection -> IO a -> IO a
sourceAction Connection
conn IO a
action = IO a -> IO a
forall a. IO a -> IO a
mask_ (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ IO Bool -> (Bool -> IO ()) -> (Bool -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
  (IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Connection -> IORef Bool
connInSource Connection
conn) (\Bool
old -> (Bool
True, Bool
old)))
  (IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
connInSource Connection
conn)) (IO a -> Bool -> IO a
forall a b. a -> b -> a
const IO a
action)

-- | Throw 'ConnectionReleased' if this handle is not the connection's
-- current one: it was released, or 'renewLease' made a newer one.
checkLease :: Connection -> IO ()
checkLease :: Connection -> IO ()
checkLease Connection
conn = do
  current <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (Connection -> IORef Int
connCurrentLease Connection
conn)
  when (current /= connLease conn) $ throwIO ConnectionReleased

-- | Make every existing handle to the connection stale.  The caller holds
-- the lock.
invalidateLease :: Connection -> IO ()
invalidateLease :: Connection -> IO ()
invalidateLease Connection
conn = IORef Int -> (Int -> (Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Connection -> IORef Int
connCurrentLease Connection
conn) (\Int
l -> (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, ()))

-- | A new handle to the connection, making every existing one stale.  Its
-- argument may be stale (that is what it is for), so it takes the mutex
-- directly, not 'withConnectionLock': it waits for a request in progress
-- on another handle rather than cutting in.
renewLease :: Connection -> IO Connection
renewLease :: Connection -> IO Connection
renewLease Connection
conn = Mutex -> IO Connection -> IO Connection
forall a. Mutex -> IO a -> IO a
withMutex (Connection -> Mutex
connMutex Connection
conn) (IO Connection -> IO Connection) -> IO Connection -> IO Connection
forall a b. (a -> b) -> a -> b
$ do
  lease <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Connection -> IORef Int
connCurrentLease Connection
conn) (\Int
l -> (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  pure conn { connLease = lease }

-- | 'renewLease', after throwing what a request would throw if the connection
-- is unusable.  The check comes first, so a refusal leaves the lease alone.
renewUsable :: Connection -> IO Connection
renewUsable :: Connection -> IO Connection
renewUsable Connection
conn = Mutex -> IO Connection -> IO Connection
forall a. Mutex -> IO a -> IO a
withMutex (Connection -> Mutex
connMutex Connection
conn) (IO Connection -> IO Connection) -> IO Connection -> IO Connection
forall a b. (a -> b) -> a -> b
$ Connection -> IO ()
refuseUnusable Connection
conn IO () -> IO Connection -> IO Connection
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Connection -> IO Connection
renewLease Connection
conn

-- | Throws 'ConnectionBroken' after a failure, or @'ConnectionClosedBy'
-- 'ClosedByClient'@ after 'closeConnection'.
checkNotBroken :: Connection -> IO ()
checkNotBroken :: Connection -> IO ()
checkNotBroken Connection
conn = do
  TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just SomeException
e | SomeException -> Bool
closedByClient SomeException
e -> SomeException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
    Maybe SomeException
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  IORef (Maybe UnusableReason) -> IO (Maybe UnusableReason)
forall a. IORef a -> IO a
readIORef (Connection -> IORef (Maybe UnusableReason)
connBroken Connection
conn) IO (Maybe UnusableReason)
-> (Maybe UnusableReason -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe UnusableReason
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just UnusableReason
r -> ConnectionBroken -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (UnusableReason -> ConnectionBroken
ConnectionBroken UnusableReason
r)

closedByClient :: SomeException -> Bool
closedByClient :: SomeException -> Bool
closedByClient SomeException
e = case SomeException -> Maybe ConnectionClosed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
  Just (ConnectionClosedBy CloseReason
ClosedByClient) -> Bool
True
  Maybe ConnectionClosed
_ -> Bool
False

-- | Throw the connection's reason if it is broken, closed by the client,
-- or ended by the server (which marks it broken too), as a request on it
-- would.  For the functions that would otherwise return normally with
-- nothing to do ('cancelActiveStatement', the transaction functions'
-- checks), so that a normal return says the connection is not known to
-- be unusable.
refuseUnusable :: Connection -> IO ()
refuseUnusable :: Connection -> IO ()
refuseUnusable Connection
conn = do
  Connection -> IO ()
checkNotBroken Connection
conn
  TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO (ZonkAny 5)) -> Maybe SomeException -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\SomeException
e -> Connection -> UnusableReason -> IO ()
markBroken Connection
conn ((SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
ReaderThreadDied SomeException
e) IO () -> IO (ZonkAny 5) -> IO (ZonkAny 5)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> IO (ZonkAny 5)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e)

-- | Why the connection is known to be unusable, as 'refuseUnusable' would
-- throw it.  It reads what is recorded and leaves it as it is.
currentUnusableReason :: Connection -> IO (Maybe UnusableReason)
currentUnusableReason :: Connection -> IO (Maybe UnusableReason)
currentUnusableReason Connection
conn = Mutex -> IO (Maybe UnusableReason) -> IO (Maybe UnusableReason)
forall a. Mutex -> IO a -> IO a
withMutex (Connection -> Mutex
connMutex Connection
conn) (IO (Maybe UnusableReason) -> IO (Maybe UnusableReason))
-> IO (Maybe UnusableReason) -> IO (Maybe UnusableReason)
forall a b. (a -> b) -> a -> b
$ do
  ended <- TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
  broken <- readIORef (connBroken conn)
  pure $ case ended of
    Just SomeException
e | SomeException -> Bool
closedByClient SomeException
e -> UnusableReason -> Maybe UnusableReason
forall a. a -> Maybe a
Just (CloseReason -> UnusableReason
Closed CloseReason
ClosedByClient)
    Maybe SomeException
_ -> Maybe UnusableReason
broken Maybe UnusableReason
-> Maybe UnusableReason -> Maybe UnusableReason
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ((SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
ReaderThreadDied (SomeException -> UnusableReason)
-> Maybe SomeException -> Maybe UnusableReason
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SomeException
ended)

-- | Record why the connection is unusable.  The first reason stays.
markBroken :: Connection -> UnusableReason -> IO ()
markBroken :: Connection -> UnusableReason -> IO ()
markBroken Connection
conn UnusableReason
reason = IORef (Maybe UnusableReason)
-> (Maybe UnusableReason -> Maybe UnusableReason) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef (Maybe UnusableReason)
connBroken Connection
conn) (Maybe UnusableReason
-> (UnusableReason -> Maybe UnusableReason)
-> Maybe UnusableReason
-> Maybe UnusableReason
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (UnusableReason -> Maybe UnusableReason
forall a. a -> Maybe a
Just UnusableReason
reason) UnusableReason -> Maybe UnusableReason
forall a. a -> Maybe a
Just)

-- | The reason an exception gives by its own type, whatever was going on; the
-- function says what to call any other exception.
reasonFor :: (SomeException -> UnusableReason) -> SomeException -> UnusableReason
reasonFor :: (SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
other SomeException
e
  -- Don't nest: a read that fails because the reader died already says so.
  | Just (ConnectionBroken UnusableReason
r) <- SomeException -> Maybe ConnectionBroken
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = UnusableReason
r
  | Just (ConnectionClosedBy CloseReason
r) <- SomeException -> Maybe ConnectionClosed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = CloseReason -> UnusableReason
Closed CloseReason
r
  | Just (ProtocolError [Char]
m) <- SomeException -> Maybe ProtocolError
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = [Char] -> UnusableReason
ServerProtocolViolation [Char]
m
  | Bool
otherwise = SomeException -> UnusableReason
other SomeException
e

-- | The reason for a failure of the request that carried a reset: the server's
-- refusal, or a request that could not be finished.
resetFailure :: SomeException -> UnusableReason
resetFailure :: SomeException -> UnusableReason
resetFailure SomeException
e
  | Just SqlServerError
err <- SomeException -> Maybe SqlServerError
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = Maybe ServerMessage -> UnusableReason
ResetRefused (ServerMessage -> Maybe ServerMessage
forall a. a -> Maybe a
Just (SqlServerError -> ServerMessage
sqlError SqlServerError
err))
  | Just (SevereError{}) <- SomeException -> Maybe SevereError
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = Maybe ServerMessage -> UnusableReason
ResetRefused Maybe ServerMessage
forall a. Maybe a
Nothing
  | Bool
otherwise = (SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
ResetFailed SomeException
e

-- | Throw 'SevereError' for a failed DONE.  After DONE_SRVERROR the server
-- closes the connection, so it is broken from here on, as after a severity 20
-- error, without waiting for the reader thread to see the close.
throwSevere :: Connection -> Done -> Maybe Text -> IO a
throwSevere :: forall a. Connection -> Done -> Maybe Text -> IO a
throwSevere Connection
conn Done
d Maybe Text
stmt = do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Done -> Bool
doneServerError Done
d) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> UnusableReason -> IO ()
markBroken Connection
conn (Done -> UnusableReason
ServerFailedConnection Done
d)
  SevereError -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (Done -> Maybe Text -> SevereError
SevereError Done
d Maybe Text
stmt)

-- | Fail the connection if the action throws anything that leaves the
-- byte stream in an unknown position.  Async exceptions are the intended
-- interruption mechanism and leave the buffers consistent, so they don't.
breakOnFailure :: Connection -> IO a -> IO a
breakOnFailure :: forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn IO a
act = IO a
act IO a -> (SomeException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \SomeException
e -> do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SomeException -> Bool
isBenign SomeException
e) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> UnusableReason -> IO ()
markBroken Connection
conn ((SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
ReadFailed SomeException
e)
  SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
  where
    isBenign :: SomeException -> Bool
isBenign SomeException
e =
      Proxy ResultAbandoned -> SomeException -> Bool
forall ex. Exception ex => Proxy ex -> SomeException -> Bool
isA (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @ResultAbandoned) SomeException
e Bool -> Bool -> Bool
|| Proxy SqlServerError -> SomeException -> Bool
forall ex. Exception ex => Proxy ex -> SomeException -> Bool
isA (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @SqlServerError) SomeException
e
        Bool -> Bool -> Bool
|| Proxy AsyncException -> SomeException -> Bool
forall ex. Exception ex => Proxy ex -> SomeException -> Bool
isA (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @AsyncException) SomeException
e Bool -> Bool -> Bool
|| Proxy SomeAsyncException -> SomeException -> Bool
forall ex. Exception ex => Proxy ex -> SomeException -> Bool
isA (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @SomeAsyncException) SomeException
e
    isA :: forall ex. Exception ex => Proxy ex -> SomeException -> Bool
    isA :: forall ex. Exception ex => Proxy ex -> SomeException -> Bool
isA Proxy ex
_ SomeException
e = case SomeException -> Maybe ex
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e :: Maybe ex of
      Just ex
_ -> Bool
True
      Maybe ex
Nothing -> Bool
False

------------------------------------------------------------------------------
-- Byte-level I/O

-- | The typed form of a synchronous exception from the transport: an
-- 'IOException' or 'TLS.TLSException' becomes 'ConnectionClosed'.  TLS
-- reports the end of the stream (after a TCP EOF without close_notify, as
-- when the server kills the session) as an error, but it is a close like any
-- other.  Anything else, async exceptions included, is returned unchanged.
transportError :: SomeException -> SomeException
transportError :: SomeException -> SomeException
transportError SomeException
e
  | Just (SomeAsyncException
_ :: SomeAsyncException) <- SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = SomeException
e
  | Just (ConnectionClosed
_ :: ConnectionClosed) <- SomeException -> Maybe ConnectionClosed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = SomeException
e
  | Just TLSException
tls <- SomeException -> Maybe TLSException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e =
      ConnectionClosed -> SomeException
forall e. Exception e => e -> SomeException
toException (ConnectionClosed -> SomeException)
-> ConnectionClosed -> SomeException
forall a b. (a -> b) -> a -> b
$ CloseReason -> ConnectionClosed
ConnectionClosedBy (CloseReason -> ConnectionClosed)
-> CloseReason -> ConnectionClosed
forall a b. (a -> b) -> a -> b
$ if TLSException -> Bool
tlsEof TLSException
tls then CloseReason
ClosedByServer else SomeException -> CloseReason
TransportFailed SomeException
e
  | Just (IOException
_ :: IOException) <- SomeException -> Maybe IOException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = ConnectionClosed -> SomeException
forall e. Exception e => e -> SomeException
toException (ConnectionClosed -> SomeException)
-> ConnectionClosed -> SomeException
forall a b. (a -> b) -> a -> b
$ CloseReason -> ConnectionClosed
ConnectionClosedBy (SomeException -> CloseReason
TransportFailed SomeException
e)
  | Bool
otherwise = SomeException
e
  where
    tlsEof :: TLSException -> Bool
tlsEof = \case
      TLS.PostHandshake TLSError
TLS.Error_EOF -> Bool
True
      TLS.Terminated Bool
_ [Char]
_ TLSError
TLS.Error_EOF -> Bool
True
      TLS.Uncontextualized TLSError
TLS.Error_EOF -> Bool
True
      TLSException
_ -> Bool
False

-- | Read from the transport directly.  Only safe before 'startReader', since an
-- async exception inside @TLS.recvData@ corrupts the TLS session.  Throws
-- 'ConnectionClosed' on EOF or a transport failure.
recvDirect :: Connection -> IO B.ByteString
recvDirect :: Connection -> IO ByteString
recvDirect Connection
conn = IORef Transport -> IO ByteString
recvTransport (Connection -> IORef Transport
connTransport Connection
conn)

recvTransport :: IORef Transport -> IO B.ByteString
recvTransport :: IORef Transport -> IO ByteString
recvTransport IORef Transport
transport = do
  bs <- IO ByteString
recv IO ByteString -> (SomeException -> IO ByteString) -> IO ByteString
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (SomeException -> IO ByteString
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (SomeException -> IO ByteString)
-> (SomeException -> SomeException)
-> SomeException
-> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> SomeException
transportError)
  when (B.null bs) $ throwIO ConnectionClosed
  pure bs
  where
    recv :: IO ByteString
recv = IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef IORef Transport
transport IO Transport -> (Transport -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      TransportPlain Backend
backend -> Backend -> IO ByteString
recvFresh Backend
backend
      TransportTls Context
ctx -> Context -> IO ByteString
forall (m :: * -> *). MonadIO m => Context -> m ByteString
TLS.recvData Context
ctx

-- | Receive into a buffer of its own: for the few reads of a plaintext
-- handshake.  Empty at the end of the stream.
recvFresh :: Backend -> IO B.ByteString
recvFresh :: Backend -> IO ByteString
recvFresh Backend
backend = do
  fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
BI.mallocByteString Int
recvBufferSize
  n <- withForeignPtr fp (\Ptr Word8
p -> Backend -> Ptr Word8 -> Int -> IO Int
backendRecv Backend
backend Ptr Word8
p Int
recvBufferSize)
  pure (BI.BS fp (max 0 n))

-- | Move transport reads onto a dedicated thread.  Consumers then block on a
-- 'TMVar', which can be interrupted without losing bytes or TLS state.  Call
-- once the transport is final (after login).
--
-- When the reader stops, it records why in 'connEnded': the
-- 'ConnectionClosed' that 'recvDirect' threw, or 'ConnectionBroken' if it
-- died of anything else (it is never interrupted except by
-- 'closeConnection', which records its own reason first, or by the
-- finalizer below).  It never passes an exception on as data, so a reader
-- killed by 'closeConnection' doesn't look to the thread waiting for a
-- response as if that thread itself had been killed.
--
-- The thread refers to the transport, 'connIncoming' and 'connEnded' only,
-- never to the whole 'Connection': it is always reachable (blocked in a
-- read), so if it kept the connection reachable, a connection dropped
-- without 'close' would never be found garbage, and the finalizer below,
-- which stops the thread and closes the socket, would never run.
--
-- Small chunks are copied into shared buffers ('copyIntoArena') before
-- they are passed on.
--
-- On a MARS connection ('enableMars' came first) the thread cuts what it
-- receives into SMP packets instead and hands each session its TDS
-- packets ('demultiplex'); it then refers to the session registry, and
-- still not to the connection.
startReader :: Connection -> IO ()
startReader :: Connection -> IO ()
startReader Connection
conn = do
  arena <- Arena -> IO (IORef Arena)
forall a. a -> IO (IORef a)
newIORef Arena
emptyArena
  let !transport = Connection -> IORef Transport
connTransport Connection
conn
      !incoming = Connection -> TMVar ByteString
connIncoming Connection
conn
      !endedVar = Connection -> TVar (Maybe SomeException)
connEnded Connection
conn
      !backend = Connection -> Backend
connBackend Connection
conn
      receive = IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef IORef Transport
transport IO Transport -> (Transport -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        TransportPlain Backend
b -> Backend -> IORef Arena -> IO ByteString
recvIntoArena Backend
b IORef Arena
arena
        TransportTls Context
_ -> IORef Transport -> IO ByteString
recvTransport IORef Transport
transport IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IORef Arena -> ByteString -> IO ByteString
copyIntoArena IORef Arena
arena
      plainLoop = do
        bs <- IO ByteString
receive
        atomically $ putTMVar incoming bs
        plainLoop
      !loop = case Connection -> Maybe Mars
connMars Connection
conn of
        Maybe Mars
Nothing -> IO b
forall {b}. IO b
plainLoop
        Just Mars{marsRegistry :: Mars -> TVar (IntMap SmpLink)
marsRegistry = !TVar (IntMap SmpLink)
registry} -> TVar (IntMap SmpLink) -> IO ByteString -> IO b
forall a. TVar (IntMap SmpLink) -> IO ByteString -> IO a
demultiplex TVar (IntMap SmpLink)
registry IO ByteString
receive
      ended SomeException
e = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        prev <- TVar (Maybe SomeException) -> STM (Maybe SomeException)
forall a. TVar a -> STM a
readTVar TVar (Maybe SomeException)
endedVar
        case prev of
          Just SomeException
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Maybe SomeException
Nothing -> TVar (Maybe SomeException) -> Maybe SomeException -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe SomeException)
endedVar (Maybe SomeException -> STM ()) -> Maybe SomeException -> STM ()
forall a b. (a -> b) -> a -> b
$ SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (SomeException -> Maybe SomeException)
-> SomeException -> Maybe SomeException
forall a b. (a -> b) -> a -> b
$ if
            | Just (ConnectionClosed
_ :: ConnectionClosed) <- SomeException -> Maybe ConnectionClosed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> SomeException
e
            -- What 'demultiplex' found wrong with the server's packets.
            | Just (ProtocolError
_ :: ProtocolError) <- SomeException -> Maybe ProtocolError
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> SomeException
e
            | Bool
otherwise -> ConnectionBroken -> SomeException
forall e. Exception e => e -> SomeException
toException (UnusableReason -> ConnectionBroken
ConnectionBroken (SomeException -> UnusableReason
ReaderThreadDied SomeException
e))
  -- Unmasked whatever the caller's masking state ('connect' may run inside
  -- 'mask' or 'uninterruptibleMask'): the thread is stopped by 'killThread'
  -- ('stopReader'), which would otherwise wait for it (forever, if
  -- uninterruptibly masked).
  --
  -- Unbound by default: wherever the scheduler puts it, and free to move.
  -- With 'connPinReader', on the capability of the thread that connects:
  -- when that is the thread that then waits for the responses, a chunk
  -- the reader passes on wakes it without a message to another
  -- capability (with @+RTS -N@ on 32 hardware threads, a local round trip took 40
  -- to 80 us less, 'Bench.Live' "round trip"), but the readers of every
  -- connection opened by one thread (a pool's) then share that
  -- capability, and they decrypt and copy what they receive.
  let fork :: ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
      fork (forall a. IO a -> IO a) -> IO ()
act
        | Connection -> Bool
connPinReader Connection
conn = do
            (cap, _) <- ThreadId -> IO (Int, Bool)
threadCapability (ThreadId -> IO (Int, Bool)) -> IO ThreadId -> IO (Int, Bool)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ThreadId
myThreadId
            forkOnWithUnmask cap act
        | Bool
otherwise = ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (forall a. IO a -> IO a) -> IO ()
act
  tid <- mask_ $ fork $ \forall a. IO a -> IO a
unmask ->
    IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (IO () -> IO ()
forall a. IO a -> IO a
unmask IO ()
forall {b}. IO b
loop) IO (Either SomeException ())
-> (Either SomeException () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO ())
-> (() -> IO ()) -> Either SomeException () -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> IO ()
ended (\() -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
  -- Named for 'GHC.Conc.listThreads' users and the eventlog.
  labelThread tid "hmssql: connection reader"
  weakTid <- mkWeakThreadId tid
  writeIORef (connReader conn) (Just weakTid)
  -- If the connection is dropped without 'close', don't leak the thread or
  -- the session.  (Closing a backend twice is harmless.)  As in
  -- 'closeConnection', the reader is stopped before the backend is closed.
  -- No thread is sending then: 'sendBytes' keeps the connection
  -- reachable until it is done, so the finalizer can't run before.
  -- The finalizer refers to the backend, never to the connection.
  _ <- mkWeakIORef (connAlive conn) (stopReader backend weakTid >> backendClose backend)
  pure ()

-- | The reader thread's loop on a MARS connection: cut the bytes received
-- into SMP packets (a header may straddle two chunks) and act on each.
-- Returns never; throws 'ConnectionClosed' when the server closes the
-- connection, or a session the client didn't close, and 'ProtocolError'
-- when a packet breaks the protocol.
--
-- A packet is handled in one STM transaction, which never blocks: the
-- server sends no further than the window it was given, so the queues
-- are bounded by it, and a packet beyond the window is an error.
demultiplex :: TVar (IntMap.IntMap SmpLink) -> IO B.ByteString -> IO a
demultiplex :: forall a. TVar (IntMap SmpLink) -> IO ByteString -> IO a
demultiplex TVar (IntMap SmpLink)
registry IO ByteString
receive = ByteString -> IO a
forall {a}. ByteString -> IO a
go ByteString
B.empty
  where
    go :: ByteString -> IO a
go ByteString
buf
      | ByteString -> Int
B.length ByteString
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
smpHeaderLength = ByteString -> IO a
more ByteString
buf
      | ByteString -> Word8
BU.unsafeHead ByteString
buf Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
smpIdentifier = [Char] -> IO a
forall a. [Char] -> IO a
violation [Char]
"a packet that is not an SMP packet"
      | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
smpHeaderLength Bool -> Bool -> Bool
|| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
smpHeaderLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
65535 = [Char] -> IO a
forall a. [Char] -> IO a
violation [Char]
"an SMP packet of an impossible length"
      | SmpHeader -> Word8
smpFlags SmpHeader
hdr Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= SmpFlag -> Word8
smpFlagWord SmpFlag
SmpData Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
smpHeaderLength = [Char] -> IO a
forall a. [Char] -> IO a
violation [Char]
"a control packet with a payload"
      | ByteString -> Int
B.length ByteString
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = ByteString -> IO a
more ByteString
buf
      | Bool
otherwise = do
          STM (Maybe SomeException) -> IO (Maybe SomeException)
forall a. STM a -> IO a
atomically (SmpHeader -> ByteString -> STM (Maybe SomeException)
packet SmpHeader
hdr (Int -> ByteString -> ByteString
BU.unsafeTake Int
len ByteString
buf)) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO (ZonkAny 3)) -> Maybe SomeException -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SomeException -> IO (ZonkAny 3)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO
          ByteString -> IO a
go (Int -> ByteString -> ByteString
BU.unsafeDrop Int
len ByteString
buf)
      where
        hdr :: SmpHeader
hdr = ByteString -> SmpHeader
readSmpHeader ByteString
buf
        len :: Int
len = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SmpHeader -> Word32
smpLength SmpHeader
hdr) :: Int

    -- A packet split between chunks is copied (as 'takeBytes' does).
    more :: ByteString -> IO a
more ByteString
buf = do
      bs <- IO ByteString
receive
      go (if B.null buf then bs else buf <> bs)

    violation :: String -> IO a
    violation :: forall a. [Char] -> IO a
violation [Char]
what = ProtocolError -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (ProtocolError -> IO a) -> ProtocolError -> IO a
forall a b. (a -> b) -> a -> b
$ [Char] -> ProtocolError
ProtocolError ([Char] -> ProtocolError) -> [Char] -> ProtocolError
forall a b. (a -> b) -> a -> b
$ [Char]
"SMP: the server sent " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
what

    wrong :: [Char] -> f (Maybe SomeException)
wrong [Char]
what = Maybe SomeException -> f (Maybe SomeException)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe SomeException -> f (Maybe SomeException))
-> Maybe SomeException -> f (Maybe SomeException)
forall a b. (a -> b) -> a -> b
$ SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (SomeException -> Maybe SomeException)
-> SomeException -> Maybe SomeException
forall a b. (a -> b) -> a -> b
$ ProtocolError -> SomeException
forall e. Exception e => e -> SomeException
toException (ProtocolError -> SomeException) -> ProtocolError -> SomeException
forall a b. (a -> b) -> a -> b
$ [Char] -> ProtocolError
ProtocolError ([Char] -> ProtocolError) -> [Char] -> ProtocolError
forall a b. (a -> b) -> a -> b
$ [Char]
"SMP: the server sent " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
what

    packet :: SmpHeader -> B.ByteString -> STM (Maybe SomeException)
    packet :: SmpHeader -> ByteString -> STM (Maybe SomeException)
packet SmpHeader
hdr ByteString
bytes = do
      links <- TVar (IntMap SmpLink) -> STM (IntMap SmpLink)
forall a. TVar a -> STM a
readTVar TVar (IntMap SmpLink)
registry
      case IntMap.lookup (fromIntegral (smpSession hdr)) links of
        Maybe SmpLink
Nothing -> [Char] -> STM (Maybe SomeException)
forall {f :: * -> *}.
Applicative f =>
[Char] -> f (Maybe SomeException)
wrong ([Char]
"a packet for session " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Word16 -> [Char]
forall a. Show a => a -> [Char]
show (SmpHeader -> Word16
smpSession SmpHeader
hdr) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
", which is not open")
        Just SmpLink
link -> do
          -- The window only grows.
          highWater <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (SmpLink -> TVar Word32
linkSendHighWater SmpLink
link)
          let grown = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SmpHeader -> Word32
smpWindow SmpHeader
hdr Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
highWater) :: Int32
          if grown < 0
            then wrong "a window smaller than its last"
            else do
              when (grown > 0) $ writeTVar (linkSendHighWater link) (smpWindow hdr)
              if
                | flags == smpFlagWord SmpData -> dataPacket link
                | flags == smpFlagWord SmpAck -> pure Nothing
                | flags == smpFlagWord SmpFin -> do
                    closing <- readTVar (linkClosing link)
                    if closing
                      then Nothing <$ writeTVar registry (IntMap.delete (fromIntegral (smpSession hdr)) links)
                      -- The server closes a session on its own only as it
                      -- closes the connection (a KILL, a severity 20 error).
                      else pure (Just (toException (ConnectionClosedBy ClosedByServer)))
                | otherwise -> wrong ("a packet with flags " ++ show flags)
      where
        flags :: Word8
flags = SmpHeader -> Word8
smpFlags SmpHeader
hdr
        tds :: ByteString
tds = Int -> ByteString -> ByteString
BU.unsafeDrop Int
smpHeaderLength ByteString
bytes
        dataPacket :: SmpLink -> STM (Maybe SomeException)
dataPacket SmpLink
link = do
          received <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (SmpLink -> TVar Word32
linkRecvSeq SmpLink
link)
          accepted <- readTVar (linkRecvHighWater link)
          if
            | smpSeqNum hdr /= received + 1 -> wrong "a DATA packet out of sequence"
            | (fromIntegral (accepted - smpSeqNum hdr) :: Int32) < 0 -> wrong "a DATA packet beyond the window"
            -- One TDS packet, whole, in each.
            | B.length tds < headerLength || fromIntegral (headerPacketLength (readHeader tds)) /= B.length tds ->
                wrong "a DATA packet that is not one TDS packet"
            | otherwise -> do
                writeTVar (linkRecvSeq link) (smpSeqNum hdr)
                when (headerStatus (readHeader tds) `hasStatus` statusEndOfMessage) $
                  modifyTVar' (linkEoms link) (+ 1)
                writeTQueue (linkIncoming link) tds
                pure Nothing

-- | Stop the reader thread, unless it has finished and been collected:
-- end its receive ('backendStopReceiving'), then kill it.  Where the
-- receive waits in GHC's I/O manager the first does nothing and the
-- exception ends the wait; where a receive can't be interrupted the first
-- ends it, and the exception finds the thread on its way out.
stopReader :: Backend -> Weak ThreadId -> IO ()
stopReader :: Backend -> Weak ThreadId -> IO ()
stopReader Backend
backend Weak ThreadId
weakTid = do
  Backend -> IO ()
backendStopReceiving Backend
backend
  Weak ThreadId -> IO (Maybe ThreadId)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak ThreadId
weakTid IO (Maybe ThreadId) -> (Maybe ThreadId -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ThreadId -> IO ()) -> Maybe ThreadId -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ThreadId -> IO ()
killThread

-- | A buffer that received chunks are copied into, and how much of it is
-- used.  Only the reader thread writes to it, and only past what it has
-- handed out, so the slices handed out never change.
data Arena = Arena !(ForeignPtr Word8) !Int

emptyArena :: Arena
emptyArena :: Arena
emptyArena = ForeignPtr Word8 -> Int -> Arena
Arena ForeignPtr Word8
nullForeignPtr Int
arenaSize
  where
    BI.BS ForeignPtr Word8
nullForeignPtr Int
_ = ByteString
B.empty

-- | Big enough that GHC allocates it as a large object, in blocks of its own.
arenaSize :: Int
arenaSize :: Int
arenaSize = Int
32768

-- | A copy of a small chunk, in an arena shared with the chunks received
-- before and after it; a chunk of 4096 bytes or more is returned unchanged.
--
-- Over TLS, @TLS.recvData@ returns a record at a time, and SQL Server sends
-- each packet as a record of its own.  With small packets, the plaintext of
-- a record is a small pinned object, allocated in the same blocks as the
-- short-lived objects tls and crypton allocate per record, many of which
-- carry a finalizer (@ScrubbedBytes@).  GHC frees a block of small pinned
-- objects only once all of them are dead, and the chunks of a large token
-- stay live until the whole token has been read ('arState').  So the chunks
-- kept their blocks, and the finalizable objects in them (about four per
-- record), alive and promoted, and those could only be found dead by a
-- major GC, which keeps them for their finalizers: after reading a 32 MiB
-- value with 512-byte packets, about 180 MiB was still live at the first
-- major GC, and freed only at the second.  Copied into an arena, the chunks
-- keep nothing else alive.  (Copying each chunk on its own would put the
-- copy in the same blocks.)  The cost: a chunk still referred to (say, the
-- rest of the last packet read) keeps its whole arena alive.
copyIntoArena :: IORef Arena -> B.ByteString -> IO B.ByteString
copyIntoArena :: IORef Arena -> ByteString -> IO ByteString
copyIntoArena IORef Arena
ref bs :: ByteString
bs@(BI.BS ForeignPtr Word8
src Int
len)
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
4096 = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
bs
  | Bool
otherwise = do
      Arena fp0 used0 <- IORef Arena -> IO Arena
forall a. IORef a -> IO a
readIORef IORef Arena
ref
      (fp, used) <-
        if used0 + len <= arenaSize
          then pure (fp0, used0)
          else (,0) <$> BI.mallocByteString arenaSize
      -- A bounded memcpy: within unsafeWithForeignPtr's contract.
      BI.unsafeWithForeignPtr fp $ \Ptr Word8
dst -> ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
BI.unsafeWithForeignPtr ForeignPtr Word8
src ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
s ->
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
used) Ptr Word8
s Int
len
      writeIORef ref (Arena fp (used + len))
      pure (BI.BS (fp `plusForeignPtr` used) len)

-- | Receive from the backend straight into the free end of the arena
-- (a new one when less than 'minRecvSpace' is left): a slice of it,
-- neither allocated nor copied per call.  (@Socket.recv sock 32768@
-- allocated 32 KiB every time, and copied what arrived when it was less:
-- most of a small request's allocation.)  Throws like 'recvTransport'.
--
-- The buffer is kept alive over the receive with 'withForeignPtr'
-- (@keepAlive#@), never @unsafeWithForeignPtr@: a read blocks
-- for as long as the server takes, can fail, and is where the reader
-- thread is interrupted, none of which the latter's contract (a
-- continuation that can't diverge) allows.
recvIntoArena :: Backend -> IORef Arena -> IO B.ByteString
recvIntoArena :: Backend -> IORef Arena -> IO ByteString
recvIntoArena Backend
backend IORef Arena
ref = do
  Arena fp0 used0 <- IORef Arena -> IO Arena
forall a. IORef a -> IO a
readIORef IORef Arena
ref
  (fp, used) <-
    if arenaSize - used0 >= minRecvSpace
      then pure (fp0, used0)
      else (,0) <$> BI.mallocByteString arenaSize
  n <- withForeignPtr fp (\Ptr Word8
p -> Backend -> Ptr Word8 -> Int -> IO Int
backendRecv Backend
backend (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
used) (Int
arenaSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
used))
    `catch` (throwIO . transportError)
  when (n <= 0) $ throwIO ConnectionClosed
  writeIORef ref (Arena fp (used + n))
  pure (BI.BS (fp `plusForeignPtr` used) n)

-- | A receive function for a TLS backend: exactly the bytes asked for, or
-- fewer at the end of the stream (as @TLS.getBackend@'s), copied out of a
-- 32 KiB buffer that is filled by receiving as much as the backend has.
-- tls reads each record in two calls, its 5-byte header and then its body,
-- which @TLS.getBackend@ made two system calls; here one receive usually
-- covers several records.  Each call allocates what it returns, as
-- @TLS.getBackend@'s did; the buffer is the connection's for its life.
-- One thread reads at a time (the connecting thread during the handshake,
-- the reader thread after).  The buffer is kept alive over each receive
-- with 'withForeignPtr', as in 'recvIntoArena'.
newBufferedRecv :: Backend -> IO (Int -> IO B.ByteString)
newBufferedRecv :: Backend -> IO (Int -> IO ByteString)
newBufferedRecv Backend
backend = do
  buf <- Int -> IO (ForeignPtr (ZonkAny 2))
forall a. Int -> IO (ForeignPtr a)
BI.mallocByteString Int
recvBufferSize
  -- The unread bytes are those from the first offset to the second.
  window <- newIORef (0 :: Int, 0 :: Int)
  let recvN Int
n = do
        (start, end) <- IORef (Int, Int) -> IO (Int, Int)
forall a. IORef a -> IO a
readIORef IORef (Int, Int)
window
        let have = Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start
        if
          | have >= n -> do
              writeIORef window (start + n, end)
              copyOut start n
          | n > recvBufferSize -> do
              -- More than the buffer holds (tls never asks for that): what
              -- it has, then straight from the backend.
              writeIORef window (0, 0)
              BI.createUptoN n $ \Ptr Word8
dst -> do
                ForeignPtr (ZonkAny 2) -> (Ptr (ZonkAny 2) -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr (ZonkAny 2)
buf ((Ptr (ZonkAny 2) -> IO ()) -> IO ())
-> (Ptr (ZonkAny 2) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr (ZonkAny 2)
src -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
dst (Ptr (ZonkAny 2)
src Ptr (ZonkAny 2) -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
start) Int
have
                let fill :: Int -> IO Int
fill Int
k
                      | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k
                      | Bool
otherwise = Backend -> Ptr Word8 -> Int -> IO Int
backendRecv Backend
backend (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
k) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) IO Int -> (Int -> IO Int) -> IO Int
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                          Int
r | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 -> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k
                            | Bool
otherwise -> Int -> IO Int
fill (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r)
                Int -> IO Int
fill Int
have
          | otherwise -> do
              -- Move the unread bytes to the front, and receive after them.
              when (start > 0 && have > 0) $ withForeignPtr buf $ \Ptr (ZonkAny 2)
p ->
                Ptr (ZonkAny 2) -> Ptr (ZonkAny 2) -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
moveBytes Ptr (ZonkAny 2)
p (Ptr (ZonkAny 2)
p Ptr (ZonkAny 2) -> Int -> Ptr (ZonkAny 2)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
start) Int
have
              writeIORef window (0, have)
              r <- withForeignPtr buf $ \Ptr (ZonkAny 2)
p ->
                Backend -> Ptr Word8 -> Int -> IO Int
backendRecv Backend
backend (Ptr (ZonkAny 2)
p Ptr (ZonkAny 2) -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
have) (Int
recvBufferSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
have)
              if r <= 0
                then writeIORef window (0, 0) >> copyOut 0 have
                else writeIORef window (0, have + r) >> recvN n
      copyOut Int
off Int
len = Int -> (Ptr Word8 -> IO ()) -> IO ByteString
BI.create Int
len ((Ptr Word8 -> IO ()) -> IO ByteString)
-> (Ptr Word8 -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dst ->
        ForeignPtr (ZonkAny 2) -> (Ptr (ZonkAny 2) -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr (ZonkAny 2)
buf ((Ptr (ZonkAny 2) -> IO ()) -> IO ())
-> (Ptr (ZonkAny 2) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr (ZonkAny 2)
src -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
dst (Ptr (ZonkAny 2)
src Ptr (ZonkAny 2) -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off) Int
len
  pure recvN

-- | The size of the buffer of 'newBufferedRecv': more than a TLS record (at most
-- 16 KiB of plaintext and its overhead).
recvBufferSize :: Int
recvBufferSize :: Int
recvBufferSize = Int
32768

-- | The least free space 'recvIntoArena' reads into.
minRecvSpace :: Int
minRecvSpace :: Int
minRecvSpace = Int
8192

-- | Mark the connection closed by the client, stop the reader thread, run
-- @goodbye@ (a TLS close_notify) unless a send is in progress, and close the
-- socket.  A thread waiting for a response or sending, and every later use,
-- then fails with @'ConnectionClosedBy' 'ClosedByClient'@.
--
-- A send in progress (or a goodbye, from another 'closeConnection') may
-- wait for the server to read, which may be never, so it is not waited for
-- as it is: the backend is shut down, which makes the send fail at once, and
-- closed only once the send has given up the send lock.  Closing a socket
-- under a thread that is sending on it would be unsafe: that thread may
-- read the descriptor just before the close and use it just after, when
-- the number may already belong to another socket.  (Or start to wait on
-- it after the close: GHC's IO manager (before 9.14) then keeps a stale
-- registration for the number, which on macOS makes closing the number's
-- next socket fail; see 'closeSocket'.)  The wait for the send lock is
-- interruptible; the backend is then shut down but not closed, until the
-- next 'closeConnection' (or the finalizer).
--
-- The backend is closed only once the reader thread has stopped: if it were
-- still reading, it could go on to read from whatever socket the next
-- connection is given with the same descriptor.
--
-- The connection is kept reachable throughout ('keepAlive'), as a send
-- keeps it: past its first lines this refers only to the fields it uses,
-- so a connection whose last handle was the caller's could otherwise be
-- found garbage half-way, and the finalizer of 'startReader' would close the
-- backend under the goodbye.  The close_notify was then lost, its send
-- failing on the closed socket; a send that had read the descriptor just
-- before the close could have written it to whatever socket was given
-- the number next.
closeConnection :: Connection -> IO () -> IO ()
closeConnection :: Connection -> IO () -> IO ()
closeConnection Connection
conn IO ()
goodbye = IORef () -> IO () -> IO ()
forall a b. a -> IO b -> IO b
keepAlive (Connection -> IORef ()
connAlive Connection
conn) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Maybe SomeException) -> Maybe SomeException -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (ConnectionClosed -> SomeException
forall e. Exception e => e -> SomeException
toException (CloseReason -> ConnectionClosed
ConnectionClosedBy CloseReason
ClosedByClient)))
  IORef (Maybe (Weak ThreadId)) -> IO (Maybe (Weak ThreadId))
forall a. IORef a -> IO a
readIORef (Connection -> IORef (Maybe (Weak ThreadId))
connReader Connection
conn) IO (Maybe (Weak ThreadId))
-> (Maybe (Weak ThreadId) -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Weak ThreadId -> IO ()) -> Maybe (Weak ThreadId) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Backend -> Weak ThreadId -> IO ()
stopReader Backend
backend)
  ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    free <- Maybe () -> Bool
forall a. Maybe a -> Bool
isJust (Maybe () -> Bool) -> IO (Maybe ()) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVar () -> IO (Maybe ())
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar ()
lock
    unless free $ do
      backendShutdown backend
      restore (takeMVar lock)
    -- A failure to say goodbye doesn't matter (the caller's @goodbye@
    -- catches it), but an asynchronous exception must reach the code that
    -- sent it; the backend is closed either way.
    ((when free (restore goodbye) `finally` backendClose backend) `finally` putMVar lock ())
  where
    lock :: MVar ()
lock = Connection -> MVar ()
connSendLock Connection
conn
    backend :: Backend
backend = Connection -> Backend
connBackend Connection
conn

-- | After 'closeConnection': drop what only a live
-- transport needs, so that a closed connection that is still referenced (by
-- a pool, a stream, a thread) keeps little alive.  A TLS context keeps its
-- record buffers and the parameters it was made with, the certificate store
-- among them: about 3 MB with the system store.  The transport becomes the
-- (closed) backend; nothing uses it any more, since the reader thread has
-- stopped and sends and reads check 'connEnded' first.  A thread that was
-- sending keeps the TLS context it had until its send fails.
--
-- The current response goes too: a response still arriving (a stream read
-- part-way) keeps the rest of the packet or chunk it stopped in, up to a
-- whole 32 KiB arena ('copyIntoArena').  Nothing reads it again, since
-- 'nextToken' and 'startRequest' check 'connEnded' first; a stream or
-- thread that still has the response keeps it as long as it does.
releaseTransport :: Connection -> IO ()
releaseTransport :: Connection -> IO ()
releaseTransport Connection
conn = do
  IORef Transport -> Transport -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Transport
connTransport Connection
conn) (Backend -> Transport
TransportPlain (Connection -> Backend
connBackend Connection
conn))
  sessions <- Connection -> IO [Session]
allSessions Connection
conn
  forM_ sessions $ \Session
sess -> do
    IORef [ByteString] -> [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Session -> IORef [ByteString]
sessRecvBuffer Session
sess) []
    IORef (Maybe ActiveResponse) -> Maybe ActiveResponse -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess) Maybe ActiveResponse
forall a. Maybe a
Nothing
    Maybe SmpSession -> (SmpSession -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Session -> Maybe SmpSession
sessSmp Session
sess) ((SmpSession -> IO ()) -> IO ()) -> (SmpSession -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \SmpSession
smp -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM [ByteString] -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM [ByteString] -> STM ()) -> STM [ByteString] -> STM ()
forall a b. (a -> b) -> a -> b
$ TQueue ByteString -> STM [ByteString]
forall a. TQueue a -> STM [a]
flushTQueue (SmpLink -> TQueue ByteString
linkIncoming (SmpSession -> SmpLink
smpLink SmpSession
smp))
  atomically $ void $ tryTakeTMVar (connIncoming conn)

-- | Read some bytes.  Throws 'ConnectionClosed' on EOF, and after
-- 'closeConnection'.
recvSome :: Connection -> IO B.ByteString
recvSome :: Connection -> IO ByteString
recvSome Connection
conn = IORef (Maybe (Weak ThreadId)) -> IO (Maybe (Weak ThreadId))
forall a. IORef a -> IO a
readIORef (Connection -> IORef (Maybe (Weak ThreadId))
connReader Connection
conn) IO (Maybe (Weak ThreadId))
-> (Maybe (Weak ThreadId) -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Maybe (Weak ThreadId)
Nothing -> Connection -> IO ByteString
recvDirect Connection
conn
  Just Weak ThreadId
_ -> STM ByteString -> IO ByteString
forall a. STM a -> IO a
atomically (STM ByteString -> IO ByteString)
-> STM ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
    ended <- TVar (Maybe SomeException) -> STM (Maybe SomeException)
forall a. TVar a -> STM a
readTVar (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
    case ended of
      Just SomeException
e | SomeException -> Bool
closedByClient SomeException
e -> SomeException -> STM ByteString
forall e a. Exception e => e -> STM a
throwSTM SomeException
e
      -- Bytes the reader received before it stopped come first.
      Maybe SomeException
_ -> TMVar ByteString -> STM ByteString
forall a. TMVar a -> STM a
takeTMVar (Connection -> TMVar ByteString
connIncoming Connection
conn) STM ByteString -> STM ByteString -> STM ByteString
forall a. STM a -> STM a -> STM a
`orElse` STM ByteString
-> (SomeException -> STM ByteString)
-> Maybe SomeException
-> STM ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe STM ByteString
forall a. STM a
retry SomeException -> STM ByteString
forall e a. Exception e => e -> STM a
throwSTM Maybe SomeException
ended

-- | Receive more bytes into the primary session's 'sessRecvBuffer' (the
-- only one read this way: an SMP session gets whole packets).  Masked from
-- taking the bytes to storing them, so an async exception can't lose them
-- in between; the wait for them stays interruptible.
recvMore :: Connection -> IO ()
recvMore :: Connection -> IO ()
recvMore Connection
conn = IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  more <- Connection -> IO ByteString
recvSome Connection
conn
  modifyIORef' (primaryBuffer conn) (++ [more])

primaryBuffer :: Connection -> IORef [B.ByteString]
primaryBuffer :: Connection -> IORef [ByteString]
primaryBuffer = Session -> IORef [ByteString]
sessRecvBuffer (Session -> IORef [ByteString])
-> (Connection -> Session) -> Connection -> IORef [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> Session
connPrimary

-- | Read exactly @n@ bytes.  Bytes stay in 'sessRecvBuffer' until all @n@ are
-- available, so an interrupted call loses nothing.
recvExactly :: Connection -> Int -> IO B.ByteString
recvExactly :: Connection -> Int -> IO ByteString
recvExactly Connection
conn Int
n = Connection -> Int -> IO ()
fillBuffer Connection
conn Int
n IO () -> IO ByteString -> IO ByteString
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Connection -> Int -> IO ByteString
takeBytes Connection
conn Int
n

-- | Receive until 'sessRecvBuffer' holds at least @n@ bytes.
fillBuffer :: Connection -> Int -> IO ()
fillBuffer :: Connection -> Int -> IO ()
fillBuffer Connection
conn Int
n = IO ()
go
  where
    go :: IO ()
go = do
      chunks <- IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef (Connection -> IORef [ByteString]
primaryBuffer Connection
conn)
      unless (sum (map B.length chunks) >= n) $ recvMore conn >> go

-- | The first @n@ bytes of 'sessRecvBuffer', which holds at least that
-- many, without consuming them: a slice of the first chunk, or a copy of
-- the pieces of chunks they span.
peekBytes :: Connection -> Int -> IO B.ByteString
peekBytes :: Connection -> Int -> IO ByteString
peekBytes Connection
conn Int
n = (ByteString, [ByteString]) -> ByteString
forall a b. (a, b) -> a
fst ((ByteString, [ByteString]) -> ByteString)
-> ([ByteString] -> (ByteString, [ByteString]))
-> [ByteString]
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [ByteString] -> (ByteString, [ByteString])
splitChunks Int
n ([ByteString] -> ByteString) -> IO [ByteString] -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef (Connection -> IORef [ByteString]
primaryBuffer Connection
conn)

-- | 'peekBytes', consuming them.
takeBytes :: Connection -> Int -> IO B.ByteString
takeBytes :: Connection -> Int -> IO ByteString
takeBytes Connection
conn Int
n = do
  (now, later) <- Int -> [ByteString] -> (ByteString, [ByteString])
splitChunks Int
n ([ByteString] -> (ByteString, [ByteString]))
-> IO [ByteString] -> IO (ByteString, [ByteString])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef (Connection -> IORef [ByteString]
primaryBuffer Connection
conn)
  writeIORef (primaryBuffer conn) later
  pure now

-- | The first @n@ bytes of the chunks (joined only if they span chunks),
-- and the chunks after them.
splitChunks :: Int -> [B.ByteString] -> (B.ByteString, [B.ByteString])
splitChunks :: Int -> [ByteString] -> (ByteString, [ByteString])
splitChunks Int
n = \case
  ByteString
c : [ByteString]
cs | ByteString -> Int
B.length ByteString
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n ->
    let (ByteString
now, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
n ByteString
c
     in (ByteString
now, if ByteString -> Bool
B.null ByteString
rest then [ByteString]
cs else ByteString
rest ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
cs)
  [ByteString]
chunks -> Int -> [ByteString] -> [ByteString] -> (ByteString, [ByteString])
go Int
n [] [ByteString]
chunks
  where
    go :: Int -> [ByteString] -> [ByteString] -> (ByteString, [ByteString])
go Int
_ [ByteString]
acc [] = ([ByteString] -> ByteString
B.concat ([ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse [ByteString]
acc), [])
    go Int
k [ByteString]
acc (ByteString
c : [ByteString]
cs)
      | ByteString -> Int
B.length ByteString
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k = Int -> [ByteString] -> [ByteString] -> (ByteString, [ByteString])
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
c) (ByteString
c ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
acc) [ByteString]
cs
      | Bool
otherwise =
          let (ByteString
now, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
k ByteString
c
           in ([ByteString] -> ByteString
B.concat ([ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse (ByteString
now ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
acc)), if ByteString -> Bool
B.null ByteString
rest then [ByteString]
cs else ByteString
rest ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
cs)

-- | Send a message, as chunks of whole packets ('splitIntoPackets',
-- 'frameSegments').  Returns normally if and only if the whole message was
-- sent; otherwise the connection is broken.
--
-- The chunks are produced as they are sent ('sendFramed' frames them
-- lazily), so the caller must have evaluated the message body in full
-- first: a lazy encoding failure (an 'error' inside a 'Put') must fail
-- before anything goes out, not cut a message short.  The send runs with asynchronous exceptions masked, so it can only be
-- interrupted while it actually blocks: waiting for the transport to take
-- the bytes (a peer that doesn't read), or for the TLS session's write
-- lock.  Any exception that escapes it, asynchronous or not, marks the
-- connection broken (the server may hold part of a message, and neither it
-- nor the TLS session can be told to forget it), shuts the socket down so
-- the reader thread ends and a later 'closeConnection' is clean, and is
-- rethrown: an interruption reaches the code that sent it, and the
-- connection then throws 'ConnectionBroken' wrapping it.
--
-- A transport failure is thrown as 'ConnectionClosed' (see
-- 'transportError'), and nothing is sent once the connection has ended: if
-- the reader thread has seen the server close the connection (as when the
-- session was killed while idle), or 'closeConnection' was called, that is
-- the reason given.
sendBytes :: Connection -> [B.ByteString] -> IO ()
sendBytes :: Connection -> [ByteString] -> IO ()
sendBytes Connection
conn [ByteString]
packets = Connection -> (Transport -> [ByteString]) -> IO ()
sendWith Connection
conn ([ByteString] -> Transport -> [ByteString]
forall a b. a -> b -> a
const [ByteString]
packets)

-- | Frame a message body and send it ('sendBytes'), in chunks that suit the
-- transport: on a plain socket, up to 64 KiB of whole packets per write;
-- over TLS, one packet per chunk, since SQL Server (2025, strict
-- encryption) resets the connection when a TLS record holds more than one
-- TDS packet (the records are gathered into writes of up to 64 KiB
-- instead, 'SendGather').  The body must be evaluated in full already.
--
-- On an SMP session each packet goes out as an SMP DATA packet
-- ('sendSmpFrames'), and the send lock is held per write, not for the
-- message; the rest is the same: masked, and any exception out of it
-- breaks the connection.
sendFramed :: Connection -> Session -> TdsContext -> PacketType -> PacketStatus -> LB.ByteString -> IO ()
sendFramed :: Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> ByteString
-> IO ()
sendFramed Connection
conn Session
sess TdsContext
ctx PacketType
ptype PacketStatus
status ByteString
body =
  Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> (PacketWriter -> IO ())
-> IO (Either SomeException ())
forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPacketsWith Connection
conn Session
sess TdsContext
ctx PacketType
ptype PacketStatus
status (\PacketWriter
writer -> PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
writer (ByteString -> Builder
BB.lazyByteString ByteString
body))
    IO (Either SomeException ())
-> (Either SomeException () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO ())
-> (() -> IO ()) -> Either SomeException () -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Connection -> SomeException -> IO ()
forall a. Connection -> SomeException -> IO a
sendFailed Connection
conn) () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Run a send (masked); if it throws, shut the backend down and fail as
-- 'sendFailed' does.
breakingOnFailure :: Connection -> IO a -> IO a
breakingOnFailure :: forall a. Connection -> IO a -> IO a
breakingOnFailure Connection
conn IO a
act = IO a
act IO a -> (SomeException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \SomeException
e -> do
  -- Harmless if the backend is closed already ('closeConnection' shuts it
  -- down before waiting for this send to give up).
  Backend -> IO ()
backendShutdown (Connection -> Backend
connBackend Connection
conn) IO () -> (IOException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
_ :: IOException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Connection -> SomeException -> IO a
forall a. Connection -> SomeException -> IO a
sendFailed Connection
conn SomeException
e

-- | Send one TDS packet as a DATA packet on an SMP session, as
-- 'sendSmpFrames' sends each.
sendSmpPacket :: Connection -> SmpSession -> PacketSink
sendSmpPacket :: Connection -> SmpSession -> PacketSink
sendSmpPacket Connection
conn SmpSession
smp ByteString
packet = do
  _ <- Connection -> SmpSession -> IO Int
awaitSmpWindow Connection
conn SmpSession
smp
  writeSmpData conn smp [packet]

-- | How many DATA packets the server's window has room for, at least one:
-- when it is used up, this waits (interruptibly, without the send lock) for
-- the server to raise it, which it does as it consumes the packets.
awaitSmpWindow :: Connection -> SmpSession -> IO Int
awaitSmpWindow :: Connection -> SmpSession -> IO Int
awaitSmpWindow Connection
conn SmpSession
smp = do
  sent <- IORef Word32 -> IO Word32
forall a. IORef a -> IO a
readIORef (SmpSession -> IORef Word32
smpSendSeq SmpSession
smp)
  atomically $ do
    readTVar (connEnded conn) >>= mapM_ throwSTM
    highWater <- readTVar (linkSendHighWater (smpLink smp))
    let room = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
highWater Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
sent) :: Int32
    when (room <= 0) retry
    pure (fromIntegral room)

-- | Write DATA packets in one write, which the caller knows the window has
-- room for.  The headers are made under the send lock ('writeSmp'), since
-- their sequence numbers and window must be in the order they are written.
writeSmpData :: Connection -> SmpSession -> [B.ByteString] -> IO ()
writeSmpData :: Connection -> SmpSession -> [ByteString] -> IO ()
writeSmpData Connection
conn SmpSession
smp [ByteString]
packets = Connection -> IO [[ByteString]] -> IO ()
writeSmp Connection
conn (IO [[ByteString]] -> IO ()) -> IO [[ByteString]] -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  sent <- IORef Word32 -> IO Word32
forall a. IORef a -> IO a
readIORef (SmpSession -> IORef Word32
smpSendSeq SmpSession
smp)
  window <- readTVarIO (linkRecvHighWater (smpLink smp))
  -- (Not @[sent + 1 ..]@, which stops at the largest number: the sequence
  -- wraps.)
  writeIORef (smpSendSeq smp) (sent + fromIntegral (length packets))
  writeIORef (smpAdvertised smp) window
  pure (zipWith (smpDataPacket (smpSid smp) window) (iterate (+ 1) (sent + 1)) packets)

-- | Send a SYN, an ACK or a FIN on an SMP session.  It never waits for
-- the server's window, and leaves the DATA sequence number as it is: an
-- ACK must go out even when nothing more may be sent, or the two sides
-- could wait for each other.  Any exception out of it breaks the
-- connection, as for any send.
sendSmpControl :: Connection -> SmpSession -> SmpFlag -> IO ()
sendSmpControl :: Connection -> SmpSession -> SmpFlag -> IO ()
sendSmpControl Connection
conn SmpSession
smp SmpFlag
flag = Connection -> IO () -> IO ()
forall a b. a -> IO b -> IO b
keepAlive Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
breakingOnFailure Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> IO [[ByteString]] -> IO ()
writeSmp Connection
conn (SmpFlag -> SmpSession -> IO [[ByteString]]
controlPacket SmpFlag
flag SmpSession
smp)

-- | A SYN, ACK or FIN for the session, made under the send lock
-- ('writeSmp'): SEQNUM 0 for a SYN, else the last DATA packet's, and the
-- window, which is recorded as advertised.
controlPacket :: SmpFlag -> SmpSession -> IO [[B.ByteString]]
controlPacket :: SmpFlag -> SmpSession -> IO [[ByteString]]
controlPacket SmpFlag
flag SmpSession
smp = do
  seqNum <- if SmpFlag
flag SmpFlag -> SmpFlag -> Bool
forall a. Eq a => a -> a -> Bool
== SmpFlag
SmpSyn then Word32 -> IO Word32
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word32
0 else IORef Word32 -> IO Word32
forall a. IORef a -> IO a
readIORef (SmpSession -> IORef Word32
smpSendSeq SmpSession
smp)
  window <- readTVarIO (linkRecvHighWater (smpLink smp))
  writeIORef (smpAdvertised smp) window
  pure [[smpControlPacket flag (smpSid smp) seqNum window]]

-- | Write SMP packets, made once the send lock is held (so that their
-- sequence numbers and windows are in the order they are written).  Each is
-- given as its pieces: a header, and a TDS packet.  In plaintext each SMP
-- packet is one write of the backend ('backendSendMany'): a socket sends the
-- pieces as they are, and a pipe, whose server reads a write at a time,
-- joins them.  TLS joins them too, because a TLS record holds one SMP packet
-- and @sendData@ makes a record per chunk.
writeSmp :: Connection -> IO [[B.ByteString]] -> IO ()
writeSmp :: Connection -> IO [[ByteString]] -> IO ()
writeSmp Connection
conn IO [[ByteString]]
packets = MVar () -> (() -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar (Connection -> MVar ()
connSendLock Connection
conn) ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \() -> do
  TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO (ZonkAny 0)) -> Maybe SomeException -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SomeException -> IO (ZonkAny 0)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO
  smpPackets <- IO [[ByteString]]
packets
  readIORef (connTransport conn) >>= \case
    TransportPlain Backend
backend -> Backend -> [[ByteString]] -> IO ()
backendSendMany Backend
backend [[ByteString]]
smpPackets
    TransportTls Context
ctx -> do
      let gather :: SendGather
gather = Connection -> SendGather
connSendGather Connection
conn
          backend :: Backend
backend = Connection -> Backend
connBackend Connection
conn
      SendGather -> IO ()
startGathering SendGather
gather
      (([ByteString] -> IO ()) -> [[ByteString]] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Context -> ByteString -> IO ()
forall (m :: * -> *). MonadIO m => Context -> ByteString -> m ()
TLS.sendData Context
ctx (ByteString -> IO ())
-> ([ByteString] -> ByteString) -> [ByteString] -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LB.fromStrict (ByteString -> ByteString)
-> ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
B.concat) [[ByteString]]
smpPackets IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> SendGather -> Backend -> IO ()
endGathering Bool
True SendGather
gather Backend
backend)
        IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`onException` Bool -> SendGather -> Backend -> IO ()
endGathering Bool
False SendGather
gather Backend
backend

sendWith :: Connection -> (Transport -> [B.ByteString]) -> IO ()
sendWith :: Connection -> (Transport -> [ByteString]) -> IO ()
sendWith Connection
conn Transport -> [ByteString]
chunksFor = Connection -> IO () -> IO ()
forall a b. a -> IO b -> IO b
keepAlive Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar () -> (() -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar (Connection -> MVar ()
connSendLock Connection
conn) ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \() -> IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  (TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO (ZonkAny 4)) -> Maybe SomeException -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SomeException -> IO (ZonkAny 4)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO) IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` Connection -> SomeException -> IO ()
forall a. Connection -> SomeException -> IO a
sendFailed Connection
conn
  transport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
connTransport Connection
conn)
  let sendAll = case Transport
transport of
        TransportPlain Backend
backend -> PacketSink -> [ByteString] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Backend -> PacketSink
backendSend Backend
backend) (Transport -> [ByteString]
chunksFor Transport
transport)
        TransportTls Context
ctx -> do
          -- One record per chunk, gathered into backend writes, and the rest
          -- written at the end of the message.
          let gather :: SendGather
gather = Connection -> SendGather
connSendGather Connection
conn
              backend :: Backend
backend = Connection -> Backend
connBackend Connection
conn
          SendGather -> IO ()
startGathering SendGather
gather
          (PacketSink -> [ByteString] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Context -> ByteString -> IO ()
forall (m :: * -> *). MonadIO m => Context -> ByteString -> m ()
TLS.sendData Context
ctx (ByteString -> IO ()) -> (ByteString -> ByteString) -> PacketSink
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LB.fromStrict) (Transport -> [ByteString]
chunksFor Transport
transport) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> SendGather -> Backend -> IO ()
endGathering Bool
True SendGather
gather Backend
backend)
            IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`onException` Bool -> SendGather -> Backend -> IO ()
endGathering Bool
False SendGather
gather Backend
backend
  breakingOnFailure conn sendAll

-- | A send failed with this exception: the connection is broken, and the
-- exception (a transport failure as 'ConnectionClosed', or the reason the
-- connection ended) rethrown.
sendFailed :: Connection -> SomeException -> IO a
sendFailed :: forall a. Connection -> SomeException -> IO a
sendFailed Connection
conn SomeException
e = do
  e' <- case SomeException -> Maybe ConnectionClosed
forall e. Exception e => SomeException -> Maybe e
fromException (SomeException -> SomeException
transportError SomeException
e) of
    Just (ConnectionClosed
closed :: ConnectionClosed) -> SomeException -> Maybe SomeException -> SomeException
forall a. a -> Maybe a -> a
fromMaybe (ConnectionClosed -> SomeException
forall e. Exception e => e -> SomeException
toException ConnectionClosed
closed) (Maybe SomeException -> SomeException)
-> IO (Maybe SomeException) -> IO SomeException
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
    Maybe ConnectionClosed
Nothing -> SomeException -> IO SomeException
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SomeException
e
  markBroken conn (reasonFor SendFailed e')
  throwIO e'

-- | Send a message whose packets are produced as they are sent: the
-- producer gets a 'PacketWriter' for a message of this packet type, and the
-- message is what it writes, ending when it returns.  For a message too long to build first (a bulk load), or one produced from an
-- effectful source.
--
-- Like 'sendBytes', it holds the send lock and runs masked, so the
-- producer is interrupted only where it blocks (a send that waits for
-- the transport, or its own source).  The packets go out as 'sendFramed'
-- sends them: up to 64 KiB of whole packets per socket write, or one
-- TLS record each, gathered.
--
-- Returns 'Right' the producer's result once every packet it sent has
-- been handed to the transport.  A synchronous exception from the
-- producer itself (not from a send) leaves the transport at a packet
-- boundary, with the message unfinished: it is returned as 'Left', and
-- the connection is not broken, for the caller to cancel the message
-- ('startStreamingRequest' does).  A failed send, or an asynchronous
-- exception (which may land inside a send), breaks the connection and is
-- rethrown, as for 'sendBytes'.
sendPackets :: forall a. Connection -> Session -> TdsContext -> PacketType -> (PacketWriter -> IO a) -> IO (Either SomeException a)
sendPackets :: forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPackets Connection
conn Session
sess TdsContext
tdsCtx PacketType
ptype = Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPacketsWith Connection
conn Session
sess TdsContext
tdsCtx PacketType
ptype (Word8 -> PacketStatus
PacketStatus Word8
0)

sendPacketsWith :: forall a. Connection -> Session -> TdsContext -> PacketType -> PacketStatus -> (PacketWriter -> IO a) -> IO (Either SomeException a)
sendPacketsWith :: forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPacketsWith Connection
conn Session
sess TdsContext
tdsCtx PacketType
ptype PacketStatus
status = Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> IO ()
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> IO ()
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPacketsTracked Connection
conn Session
sess TdsContext
tdsCtx PacketType
ptype PacketStatus
status (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

sendPacketsTracked :: forall a. Connection -> Session -> TdsContext -> PacketType -> PacketStatus -> IO () -> (PacketWriter -> IO a) -> IO (Either SomeException a)
sendPacketsTracked :: forall a.
Connection
-> Session
-> TdsContext
-> PacketType
-> PacketStatus
-> IO ()
-> (PacketWriter -> IO a)
-> IO (Either SomeException a)
sendPacketsTracked Connection
conn Session
sess TdsContext
tdsCtx PacketType
ptype PacketStatus
status IO ()
onPacket PacketWriter -> IO a
produce = case Session -> Maybe SmpSession
sessSmp Session
sess of
  Maybe SmpSession
Nothing -> Connection
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. a -> IO b -> IO b
keepAlive Connection
conn (IO (Either SomeException a) -> IO (Either SomeException a))
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ MVar ()
-> (() -> IO (Either SomeException a))
-> IO (Either SomeException a)
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar (Connection -> MVar ()
connSendLock Connection
conn) ((() -> IO (Either SomeException a))
 -> IO (Either SomeException a))
-> (() -> IO (Either SomeException a))
-> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ \() -> IO (Either SomeException a) -> IO (Either SomeException a)
forall a. IO a -> IO a
mask_ (IO (Either SomeException a) -> IO (Either SomeException a))
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ do
    (TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn) IO (Maybe SomeException) -> (Maybe SomeException -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (SomeException -> IO (ZonkAny 6)) -> Maybe SomeException -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SomeException -> IO (ZonkAny 6)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO) IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` Connection -> SomeException -> IO ()
forall a. Connection -> SomeException -> IO a
sendFailed Connection
conn
    transport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
connTransport Connection
conn)
    -- A failure inside a send is told from one in the producer by wrapping
    -- it; the producer's own synchronous exceptions are caught here.
    let run = case Transport
transport of
          TransportPlain Backend
backend -> do
            pending <- ([[ByteString]], Int) -> IO (IORef ([[ByteString]], Int))
forall a. a -> IO (IORef a)
newIORef ([], Int
0 :: Int)
            let send ByteString
p = do
                  (ps, n) <- IORef ([[ByteString]], Int) -> IO ([[ByteString]], Int)
forall a. IORef a -> IO a
readIORef IORef ([[ByteString]], Int)
pending
                  let n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
B.length ByteString
p
                  if n' >= gatherLimit
                    then do
                      writeIORef pending ([], 0)
                      backendSendMany backend (reverse ([p] : ps)) `catch` sendFailure
                    else writeIORef pending ([p] : ps, n')
                flush = do
                  (ps, _) <- IORef ([[ByteString]], Int) -> IO ([[ByteString]], Int)
forall a. IORef a -> IO a
readIORef IORef ([[ByteString]], Int)
pending
                  writeIORef pending ([], 0)
                  unless (null ps) $ backendSendMany backend (reverse ps) `catch` sendFailure
            r <- writing send
            flush
            pure r
          TransportTls Context
ctx -> do
            let gather :: SendGather
gather = Connection -> SendGather
connSendGather Connection
conn
                backend :: Backend
backend = Connection -> Backend
connBackend Connection
conn
            SendGather -> IO ()
startGathering SendGather
gather
            (do r <- PacketSink -> IO (Either SomeException a)
writing (\ByteString
p -> Context -> ByteString -> IO ()
forall (m :: * -> *). MonadIO m => Context -> ByteString -> m ()
TLS.sendData Context
ctx (ByteString -> ByteString
LB.fromStrict ByteString
p) IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` SomeException -> IO ()
forall b. SomeException -> IO b
sendFailure)
                endGathering True gather backend `catch` sendFailure
                pure r)
              IO (Either SomeException a) -> IO () -> IO (Either SomeException a)
forall a b. IO a -> IO b -> IO a
`onException` Bool -> SendGather -> Backend -> IO ()
endGathering Bool
False SendGather
gather Backend
backend
    breakingOnFailure conn (run `catch` (throwIO . unwrap))
  -- Each packet is written as it is produced (the wait for the server's
  -- window comes between them), with the send lock held for the write only.
  Just SmpSession
smp -> Connection
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. a -> IO b -> IO b
keepAlive Connection
conn (IO (Either SomeException a) -> IO (Either SomeException a))
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ IO (Either SomeException a) -> IO (Either SomeException a)
forall a. IO a -> IO a
mask_ (IO (Either SomeException a) -> IO (Either SomeException a))
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ Connection
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a. Connection -> IO a -> IO a
breakingOnFailure Connection
conn (IO (Either SomeException a) -> IO (Either SomeException a))
-> IO (Either SomeException a) -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$
    PacketSink -> IO (Either SomeException a)
writing (\ByteString
p -> Connection -> SmpSession -> PacketSink
sendSmpPacket Connection
conn SmpSession
smp ByteString
p IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` SomeException -> IO ()
forall b. SomeException -> IO b
sendFailure) IO (Either SomeException a)
-> (SomeException -> IO (Either SomeException a))
-> IO (Either SomeException a)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (SomeException -> IO (Either SomeException a)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (SomeException -> IO (Either SomeException a))
-> (SomeException -> SomeException)
-> SomeException
-> IO (Either SomeException a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> SomeException
unwrap)
  where
    writing :: PacketSink -> IO (Either SomeException a)
    writing :: PacketSink -> IO (Either SomeException a)
writing PacketSink
sink = IO a -> IO (Either SomeException a)
forall {b}. IO b -> IO (Either SomeException b)
produced (TdsContext
-> PacketType
-> PacketStatus
-> PacketSink
-> (PacketWriter -> IO a)
-> IO a
forall a.
TdsContext
-> PacketType
-> PacketStatus
-> PacketSink
-> (PacketWriter -> IO a)
-> IO a
runPacketWriterWith TdsContext
tdsCtx PacketType
ptype PacketStatus
status (\ByteString
packet -> IO ()
onPacket IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> PacketSink
sink ByteString
packet) (PacketWriter -> IO a
produce (PacketWriter -> IO a)
-> (PacketWriter -> PacketWriter) -> PacketWriter -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. IO a -> IO a) -> PacketWriter -> PacketWriter
withSourceScope IO a -> IO a
forall a. IO a -> IO a
sourceScope))
    sourceScope :: IO b -> IO b
    sourceScope :: forall a. IO a -> IO a
sourceScope = Connection -> IO b -> IO b
forall a. Connection -> IO a -> IO a
sourceAction Connection
conn
    -- The producer's own synchronous exceptions, as a value.
    produced :: IO b -> IO (Either SomeException b)
produced IO b
act = IO b -> IO (Either SomeException b)
forall e a. Exception e => IO a -> IO (Either e a)
try IO b
act IO (Either SomeException b)
-> (Either SomeException b -> IO (Either SomeException b))
-> IO (Either SomeException b)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Right b
r -> Either SomeException b -> IO (Either SomeException b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Either SomeException b
forall a b. b -> Either a b
Right b
r)
      Left SomeException
e
        | Just (PacketSendFailed SomeException
inner) <- SomeException -> Maybe PacketSendFailed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> PacketSendFailed -> IO (Either SomeException b)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (SomeException -> PacketSendFailed
PacketSendFailed SomeException
inner)
        | Just (SomeAsyncException
_ :: SomeAsyncException) <- SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e -> SomeException -> IO (Either SomeException b)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
        | Bool
otherwise -> Either SomeException b -> IO (Either SomeException b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SomeException -> Either SomeException b
forall a b. a -> Either a b
Left SomeException
e)
    unwrap :: SomeException -> SomeException
unwrap SomeException
e = case SomeException -> Maybe PacketSendFailed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
      Just (PacketSendFailed SomeException
inner) -> SomeException
inner
      Maybe PacketSendFailed
Nothing -> SomeException
e
    sendFailure :: SomeException -> IO b
    sendFailure :: forall b. SomeException -> IO b
sendFailure = PacketSendFailed -> IO b
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (PacketSendFailed -> IO b)
-> (SomeException -> PacketSendFailed) -> SomeException -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> PacketSendFailed
PacketSendFailed

-- | A transport send failed with this exception ('sendPackets').
newtype PacketSendFailed = PacketSendFailed SomeException

instance Show PacketSendFailed where
  showsPrec :: Int -> PacketSendFailed -> ShowS
showsPrec Int
p (PacketSendFailed SomeException
a) = Int -> Addr# -> SomeException -> ShowS
forall a. Show a => Int -> Addr# -> a -> ShowS
showCon1 Int
p Addr#
"PacketSendFailed"# SomeException
a
  {-# OPAQUE showsPrec #-}

instance Exception PacketSendFailed

-- | Where the TLS records of a message go: a TLS backend's send
-- ('gatheredSend') writes each record to the backend at once, unless
-- 'sendWith' is sending a message, when it gathers them into writes of up
-- to 'gatherLimit' bytes.  Records still hold one TDS packet each
-- ('sendFramed'); only the system calls change (a 4 MB parameter took 525
-- @send@s of one 8 KB record, which cost the client more CPU than tls's
-- encryption of it).  Every backend send goes through the MVar, so records
-- tls sends on its own while a message is being gathered (the reader
-- thread's answer to a KeyUpdate, an alert) are gathered too, in order.
newtype SendGather = SendGather (MVar Gathered)

data Gathered
  = NotGathering
  | Gathering ![[B.ByteString]] !Int
    -- ^ Records not yet written, newest first, each a write of one piece
    -- ('backendSendMany'), and their total length.

newSendGather :: IO SendGather
newSendGather :: IO SendGather
newSendGather = MVar Gathered -> SendGather
SendGather (MVar Gathered -> SendGather)
-> IO (MVar Gathered) -> IO SendGather
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gathered -> IO (MVar Gathered)
forall a. a -> IO (MVar a)
newMVar Gathered
NotGathering

-- | Backend writes of gathered records are at most this long, unless one
-- record is longer (as for plaintext messages, 'sendFramed').
gatherLimit :: Int
gatherLimit :: Int
gatherLimit = Int
65536

-- | A TLS backend's send over the connection's backend: at once, or
-- gathered while 'sendWith' is sending a message.  It refers to the
-- backend and the gather only, never to the connection (see 'startReader').
gatheredSend :: SendGather -> Backend -> B.ByteString -> IO ()
gatheredSend :: SendGather -> Backend -> PacketSink
gatheredSend (SendGather MVar Gathered
v) Backend
backend ByteString
bs = MVar Gathered -> (Gathered -> IO Gathered) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar Gathered
v ((Gathered -> IO Gathered) -> IO ())
-> (Gathered -> IO Gathered) -> IO ()
forall a b. (a -> b) -> a -> b
$ \case
  Gathered
NotGathering -> Gathered
NotGathering Gathered -> IO () -> IO Gathered
forall a b. a -> IO b -> IO a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Backend -> PacketSink
backendSend Backend
backend ByteString
bs
  Gathering [[ByteString]]
pending Int
size
    | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
gatherLimit -> do
        Backend -> [[ByteString]] -> IO ()
backendSendMany Backend
backend ([[ByteString]] -> [[ByteString]]
forall a. [a] -> [a]
reverse [[ByteString]]
pending)
        Gathered -> IO Gathered
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Gathered -> IO Gathered) -> Gathered -> IO Gathered
forall a b. (a -> b) -> a -> b
$! [[ByteString]] -> Int -> Gathered
Gathering [[ByteString
bs]] Int
len
    | Bool
otherwise -> Gathered -> IO Gathered
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Gathered -> IO Gathered) -> Gathered -> IO Gathered
forall a b. (a -> b) -> a -> b
$! [[ByteString]] -> Int -> Gathered
Gathering ([ByteString
bs] [ByteString] -> [[ByteString]] -> [[ByteString]]
forall a. a -> [a] -> [a]
: [[ByteString]]
pending) (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
  where
    len :: Int
len = ByteString -> Int
B.length ByteString
bs

startGathering :: SendGather -> IO ()
startGathering :: SendGather -> IO ()
startGathering (SendGather MVar Gathered
v) = MVar Gathered -> (Gathered -> IO Gathered) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar Gathered
v (\Gathered
_ -> Gathered -> IO Gathered
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[ByteString]] -> Int -> Gathered
Gathering [] Int
0))

-- | Stop gathering, and write what is gathered if asked to (at the end of
-- a message) or drop it (after a failed send: the connection is broken).
endGathering :: Bool -> SendGather -> Backend -> IO ()
endGathering :: Bool -> SendGather -> Backend -> IO ()
endGathering Bool
write (SendGather MVar Gathered
v) Backend
backend = IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  gathered <- MVar Gathered -> IO Gathered
forall a. MVar a -> IO a
takeMVar MVar Gathered
v
  let pending = case Gathered
gathered of
        Gathering [[ByteString]]
p Int
_ | Bool
write -> [[ByteString]]
p
        Gathered
_ -> []
  unless (null pending) (backendSendMany backend (reverse pending)) `finally` putMVar v NotGathering

-- | Run the action with the value kept reachable until it is done.
-- 'sendBytes' keeps its connection so, for the finalizer of 'startReader': a
-- thread that has nothing left to do with the connection once the send
-- is over keeps it reachable only through what 'sendBytes' itself refers
-- to, which is up to the optimizer; if it didn't, the finalizer could close
-- the backend under the send.
keepAlive :: a -> IO b -> IO b
keepAlive :: forall a b. a -> IO b -> IO b
keepAlive a
x (IO State# RealWorld -> (# State# RealWorld, b #)
act) = (State# RealWorld -> (# State# RealWorld, b #)) -> IO b
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> a
-> State# RealWorld
-> (State# RealWorld -> (# State# RealWorld, b #))
-> (# State# RealWorld, b #)
forall a d b. a -> State# d -> (State# d -> b) -> b
keepAlive# a
x State# RealWorld
s State# RealWorld -> (# State# RealWorld, b #)
act)

-- | Read one packet.  Interrupted while waiting for it, it loses nothing;
-- the caller must run it masked if the packet must not be lost once read
-- (as 'nextToken' does).
--
-- This reads the primary session ('readSessionPacket'): the connection's
-- one, or SMP session 0.
readPacket :: Connection -> IO (Header, B.ByteString)
readPacket :: Connection -> IO (Header, ByteString)
readPacket Connection
conn = Connection -> Session -> IO (Header, ByteString)
readSessionPacket Connection
conn (Connection -> Session
connPrimary Connection
conn)

-- | 'readPacket' from a plain TDS stream: what a connection is before and
-- during its login, and what one without MARS stays.
readPlainPacket :: Connection -> IO (Header, B.ByteString)
readPlainPacket :: Connection -> IO (Header, ByteString)
readPlainPacket Connection
conn = do
  -- Peek at the header without consuming so a partial read can be retried.
  Connection -> Int -> IO ()
fillBuffer Connection
conn Int
headerLength
  hdr <- ByteString -> Header
readHeader (ByteString -> Header) -> IO ByteString -> IO Header
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> Int -> IO ByteString
peekBytes Connection
conn Int
headerLength
  let total = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Header -> Word16
headerPacketLength Header
hdr)
  when (total < headerLength) $ throwIO $ ProtocolError "packet length smaller than header"
  fillBuffer conn total
  packet <- takeBytes conn total
  pure (hdr, B.drop headerLength packet)

-- | Read one packet of the session's: 'readPlainPacket', or on an SMP session
-- the next packet the reader thread queued for it.  The server is told
-- of the room that makes in the window (an ACK) once it has grown by
-- the cadence, @window / 16@ packets.  Interrupted while waiting for a
-- packet, it loses nothing; the ACK's send can't be interrupted, and its
-- failure breaks the connection, as for any send, unless the connection
-- has ended ('sendAck'): the packets queued before the end are read first.
--
-- The connection is kept reachable while it waits ('keepAlive'), as
-- 'sendBytes' keeps it: the wait itself refers only to what the reader
-- thread shares, and the finalizer of 'startReader' must not run under it.
readSessionPacket :: Connection -> Session -> IO (Header, B.ByteString)
readSessionPacket :: Connection -> Session -> IO (Header, ByteString)
readSessionPacket Connection
conn Session
sess = IORef () -> IO (Header, ByteString) -> IO (Header, ByteString)
forall a b. a -> IO b -> IO b
keepAlive (Connection -> IORef ()
connAlive Connection
conn) (IO (Header, ByteString) -> IO (Header, ByteString))
-> IO (Header, ByteString) -> IO (Header, ByteString)
forall a b. (a -> b) -> a -> b
$ case Session -> Maybe SmpSession
sessSmp Session
sess of
  Maybe SmpSession
Nothing -> Connection -> IO (Header, ByteString)
readPlainPacket Connection
conn
  Just SmpSession
smp -> do
    let link :: SmpLink
link = SmpSession -> SmpLink
smpLink SmpSession
smp
    packet <- STM ByteString -> IO ByteString
forall a. STM a -> IO a
atomically (STM ByteString -> IO ByteString)
-> STM ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
      ended <- TVar (Maybe SomeException) -> STM (Maybe SomeException)
forall a. TVar a -> STM a
readTVar (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
      case ended of
        Just SomeException
e | SomeException -> Bool
closedByClient SomeException
e -> SomeException -> STM ByteString
forall e a. Exception e => e -> STM a
throwSTM SomeException
e
        -- Packets the reader received before it stopped come first.
        Maybe SomeException
_ -> TQueue ByteString -> STM ByteString
forall a. TQueue a -> STM a
readTQueue (SmpLink -> TQueue ByteString
linkIncoming SmpLink
link) STM ByteString -> STM ByteString -> STM ByteString
forall a. STM a -> STM a -> STM a
`orElse` STM ByteString
-> (SomeException -> STM ByteString)
-> Maybe SomeException
-> STM ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe STM ByteString
forall a. STM a
retry SomeException -> STM ByteString
forall e a. Exception e => e -> STM a
throwSTM Maybe SomeException
ended
    accepted <- atomically $ stateTVar (linkRecvHighWater link) (\Word32
w -> (Word32
w Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1, Word32
w Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1))
    advertised <- readIORef (smpAdvertised smp)
    let window = Word32 -> (Mars -> Word32) -> Maybe Mars -> Word32
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Word32
0 Mars -> Word32
marsWindow (Connection -> Maybe Mars
connMars Connection
conn)
    -- Uninterruptibly: this is a read, and an interruption of a read must
    -- leave the connection usable, which an interrupted send doesn't.  The
    -- ACK is 16 bytes and waits only briefly (for the TLS session's state,
    -- which the reader thread shares, or for room in the socket's buffer,
    -- which is full only while a request of this connection is still
    -- going out: a bulk load whose rows come from this stream);
    -- 'closeConnection' from another thread ends the wait.
    when (accepted - advertised >= max 1 (window `div` 16)) $
      uninterruptibleMask_ (sendAck conn smp)
    pure (readHeader packet, B.drop headerLength packet)

-- | Send an ACK ('sendSmpControl'), except on a connection that has ended
-- ('connEnded': the reader thread saw the server close it, or break the
-- protocol, or 'closeConnection' was called).  The packets the reader
-- queued before the end are still to be read, a severity-20 error among
-- them, and the server needs no more window: the ACK is skipped, and a
-- send that fails because the end came meanwhile is ignored, where
-- breaking the connection would make those packets unreadable and lose
-- the one being read.  Any other failure breaks the connection, as for
-- any send.
sendAck :: Connection -> SmpSession -> IO ()
sendAck :: Connection -> SmpSession -> IO ()
sendAck Connection
conn SmpSession
smp = Connection -> IO () -> IO ()
forall a b. a -> IO b -> IO b
keepAlive Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  r <- IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (Connection -> IO [[ByteString]] -> IO ()
writeSmp Connection
conn (SmpFlag -> SmpSession -> IO [[ByteString]]
controlPacket SmpFlag
SmpAck SmpSession
smp))
  case r of
    Right () -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Left (SomeException
e :: SomeException) -> do
      ended <- Maybe SomeException -> Bool
forall a. Maybe a -> Bool
isJust (Maybe SomeException -> Bool)
-> IO (Maybe SomeException) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
      unless ended $ breakingOnFailure conn (throwIO e)

-- | A message body.  binary's 'Put.runPut' starts with a 4 KiB buffer,
-- which GHC allocates as a large object (taking the storage manager's
-- lock) for every request, though most bodies are a few dozen bytes; this
-- starts with 512 bytes and continues in 32 KiB chunks.  The chunks aren't
-- trimmed: the body is copied into packets anyway ('frameSegments').
runBody :: Put -> LB.ByteString
runBody :: Put -> ByteString
runBody = AllocationStrategy -> ByteString -> Builder -> ByteString
BB.toLazyByteStringWith (Int -> Int -> AllocationStrategy
BB.untrimmedStrategy Int
512 Int
32752) ByteString
LB.empty (Builder -> ByteString) -> (Put -> Builder) -> Put -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Put -> Builder
forall a. PutM a -> Builder
Put.execPut

-- | Frame and send a complete message.
sendMessage :: Connection -> PacketType -> Put -> IO ()
sendMessage :: Connection -> PacketType -> Put -> IO ()
sendMessage Connection
conn PacketType
ptype Put
body = do
  ctx <- IORef TdsContext -> IO TdsContext
forall a. IORef a -> IO a
readIORef (Connection -> IORef TdsContext
connContext Connection
conn)
  let bytes = Put -> ByteString
runBody Put
body
  _ <- evaluate (LB.length bytes)
  sendFramed conn (connPrimary conn) ctx ptype (PacketStatus 0) bytes

-- | Read every packet of one message and return the concatenated body.
readWholeMessage :: Connection -> IO (PacketType, B.ByteString)
readWholeMessage :: Connection -> IO (PacketType, ByteString)
readWholeMessage Connection
conn = Connection -> Int -> IO (PacketType, ByteString)
readWholeMessageUpTo Connection
conn Int
forall a. Bounded a => a
maxBound

-- | 'readWholeMessage', failing with a 'ProtocolError' once the body is
-- longer than the given number of bytes, so a server can't make the client
-- buffer without bound.  Used for PRELOGIN, whose response is not a token
-- stream.
readWholeMessageUpTo :: Connection -> Int -> IO (PacketType, B.ByteString)
readWholeMessageUpTo :: Connection -> Int -> IO (PacketType, ByteString)
readWholeMessageUpTo Connection
conn Int
limit = Connection
-> IO (PacketType, ByteString) -> IO (PacketType, ByteString)
forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn (IO (PacketType, ByteString) -> IO (PacketType, ByteString))
-> IO (PacketType, ByteString) -> IO (PacketType, ByteString)
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString] -> IO (PacketType, ByteString)
go Int
0 []
  where
    go :: Int -> [ByteString] -> IO (PacketType, ByteString)
go !Int
size [ByteString]
acc = do
      (hdr, body) <- Connection -> IO (Header, ByteString)
readPacket Connection
conn
      let acc' = ByteString
body ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
acc
          size' = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
B.length ByteString
body
      when (size' > limit) $ throwIO $ ProtocolError $
        "the server's message is longer than " ++ show limit ++ " bytes, the most expected"
      if headerStatus hdr `hasStatus` statusEndOfMessage
        then pure (headerType hdr, B.concat (reverse acc'))
        else go size' acc'

hasStatus :: PacketStatus -> PacketStatus -> Bool
hasStatus :: PacketStatus -> PacketStatus -> Bool
hasStatus PacketStatus
a PacketStatus
b = PacketStatus
a PacketStatus -> PacketStatus -> PacketStatus
forall a. Bits a => a -> a -> a
.&. PacketStatus
b PacketStatus -> PacketStatus -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8 -> PacketStatus
PacketStatus Word8
0

------------------------------------------------------------------------------
-- Requests and responses

-- | Send a request and start reading its response.  Any unfinished previous
-- response is canceled and drained first; on a MARS connection the
-- request takes a session of its own instead, and the other responses
-- stay open ('prepareForRequest').  A pending 'requestReset' is carried
-- out first (see there).
startRequest :: Connection -> PacketType -> (TdsContext -> LB.ByteString) -> IO ActiveResponse
startRequest :: Connection
-> PacketType -> (TdsContext -> ByteString) -> IO ActiveResponse
startRequest Connection
conn PacketType
ptype TdsContext -> ByteString
body = Connection
-> Maybe Text
-> PacketType
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
startRequestFor Connection
conn Maybe Text
forall a. Maybe a
Nothing PacketType
ptype ((TdsContext -> ByteString) -> TdsContext -> PacketWriter -> IO ()
bytesBody TdsContext -> ByteString
body)

-- | A request body already in memory, as the request writer takes it.  The
-- body's length is forced before anything is written, so a lazy encoding
-- failure fails the request before any of it is sent.
bytesBody :: (TdsContext -> LB.ByteString) -> TdsContext -> PacketWriter -> IO ()
bytesBody :: (TdsContext -> ByteString) -> TdsContext -> PacketWriter -> IO ()
bytesBody TdsContext -> ByteString
body TdsContext
ctx PacketWriter
writer = do
  let bytes :: ByteString
bytes = TdsContext -> ByteString
body TdsContext
ctx
  _ <- Int64 -> IO Int64
forall a. a -> IO a
evaluate (ByteString -> Int64
LB.length ByteString
bytes)
  writeBuilder writer (BB.lazyByteString bytes)

-- | 'startRequest', with the request's SQL text: what names the response
-- if it is found unfinished ('busyStatements').
startRequestFor :: Connection -> Maybe Text -> PacketType -> (TdsContext -> PacketWriter -> IO ()) -> IO ActiveResponse
startRequestFor :: Connection
-> Maybe Text
-> PacketType
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
startRequestFor Connection
conn Maybe Text
stmt PacketType
ptype TdsContext -> PacketWriter -> IO ()
mkBody = Connection -> IO ActiveResponse -> IO ActiveResponse
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO ActiveResponse -> IO ActiveResponse)
-> IO ActiveResponse -> IO ActiveResponse
forall a b. (a -> b) -> a -> b
$ do
  IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
connInSource Connection
conn) IO Bool -> (Bool -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
active -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
active (ReentrantSource -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ReentrantSource
ReentrantSource)
  Connection -> IO ()
checkNotBroken Connection
conn
  Connection -> IO ()
prepareForRequest Connection
conn
  let resettable :: Bool
resettable = PacketType
ptype PacketType -> [PacketType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PacketType
PTSqlBatch, PacketType
PTRpc, PacketType
PTTransactionManager]
  pending <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
connResetPending Connection
conn)
  when (resettable && pending) $ do
    -- The server refuses a reset while another request is active on the
    -- connection (and then fails every request until that one is
    -- canceled), so every open response is settled first.
    forM_ (connMars conn) $ \Mars
_ -> Connection -> IO ()
finishResponse Connection
conn
    unless (null (connSessionSetup conn)) $ resetWithSessionSetup conn
  sendRequest conn stmt ptype resettable mkBody

-- | What precedes every request, with the lock held.
--
-- Without MARS: the unfinished response, if any, is canceled and drained
-- ('finishSession').
--
-- With MARS, other responses stay open, except that:
--
-- * a response canceled by an interrupted read (an Attention was sent,
--   its acknowledgement not yet read) is drained, as without MARS;
-- * inside 'withReservedSession', the reserved session's response is
--   finished, as without MARS;
-- * a request that carried RESETCONNECTION must have been answered
--   ('awaitReset');
-- * idle sessions beyond 'maxIdleSessions' are closed.
prepareForRequest :: Connection -> IO ()
prepareForRequest :: Connection -> IO ()
prepareForRequest Connection
conn = case Connection -> Maybe Mars
connMars Connection
conn of
  Maybe Mars
Nothing -> Connection -> Session -> IO ()
finishSession Connection
conn (Connection -> Session
connPrimary Connection
conn)
  Just Mars
mars -> do
    Connection -> IO ()
drainCancelled Connection
conn
    IORef (Maybe Session) -> IO (Maybe Session)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars) IO (Maybe Session) -> (Maybe Session -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Session -> IO ()) -> Maybe Session -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Connection -> Session -> IO ()
finishSession Connection
conn)
    Connection -> Mars -> IO ()
awaitReset Connection
conn Mars
mars
    Connection -> Mars -> IO ()
closeIdleSessions Connection
conn Mars
mars

-- | Drain every response that was canceled (an Attention sent, by an
-- interrupted read) and not yet read to the acknowledgement, or
-- abandoned by the function that was reading it ('readingToEnd').  The
-- caller holds the lock.
drainCancelled :: Connection -> IO ()
drainCancelled :: Connection -> IO ()
drainCancelled Connection
conn = do
  sessions <- Connection -> IO [Session]
allSessions Connection
conn
  forM_ sessions $ \Session
sess -> IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ActiveResponse -> IO ()) -> Maybe ActiveResponse -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\ActiveResponse
ar -> do
    canceled <- (Cancellation -> Cancellation -> Bool
forall a. Eq a => a -> a -> Bool
/= Cancellation
Running) (Cancellation -> Bool) -> IO Cancellation -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef Cancellation -> IO Cancellation
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef Cancellation
arCancellation ActiveResponse
ar)
    when canceled $ finishSession conn sess)

-- | Run an action that reads the response to its end, on a plain or a
-- MARS connection alike.  If it throws part-way, the response is not
-- left for the next request to find: a synchronous exception cancels and
-- drains it before it propagates ('cancelActiveResponse'); an
-- asynchronous one abandons it ('abandonResponse'), which cancels the
-- statement at once and leaves the drain to the next request (or
-- transaction function, 'release' or 'cancelActiveStatement'), so that an
-- interruption never waits for the server.  (An interruption needn't
-- land in a read: it can come just after the request was sent, or between
-- two reads.)
readingToEnd :: Connection -> ActiveResponse -> IO a -> IO a
readingToEnd :: forall a. Connection -> ActiveResponse -> IO a -> IO a
readingToEnd Connection
conn ActiveResponse
ar IO a
act = IO a
act IO a -> (SomeException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
e :: SomeException) -> do
  case SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
    Just (SomeAsyncException
_ :: SomeAsyncException) -> Connection -> ActiveResponse -> IO ()
abandonResponse Connection
conn ActiveResponse
ar
    -- A failure to cancel leaves the first exception to report (an
    -- interruption of the cancellation is reported instead).
    Maybe SomeAsyncException
Nothing -> Connection -> ActiveResponse -> IO ()
cancelActiveResponse Connection
conn ActiveResponse
ar IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
c :: SomeException) -> case SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
c of
      Just (SomeAsyncException
_ :: SomeAsyncException) -> SomeException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
c
      Maybe SomeAsyncException
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e

-- | After an asynchronous exception out of a function or scope that was
-- reading the response ('readingToEnd', 'withQueryS'): mark the response
-- abandoned, so that the next request (or transaction function, 'release'
-- or 'cancelActiveStatement') drains it; abandon the handle (a new
-- generation), so that a stream read on through it throws
-- 'ResultAbandoned' instead of ending at the acknowledgement, silently
-- missing rows; and cancel the statement at once ('cancelStatement').
-- Nothing if the response has ended or was superseded.
--
-- Runs masked (from a handler).  It never waits for the statement; it
-- may wait for the connection mutex (taken directly: the handle may be
-- stale, and the response is to be marked all the same) and, on a MARS
-- connection, for a slot in the server's window.  The abandoned flag is
-- written before the mutex is taken, so that a second interruption in
-- that wait still leaves the response for the next drain to find (on a
-- superseded response, which nothing consults, the flag is harmless); the
-- other state changes come before the send, since they can't fail.
abandonResponse :: Connection -> ActiveResponse -> IO ()
abandonResponse :: Connection -> ActiveResponse -> IO ()
abandonResponse Connection
conn ActiveResponse
ar = do
  ActiveResponse -> IO ()
markAbandoned ActiveResponse
ar
  Mutex -> IO () -> IO ()
forall a. Mutex -> IO a -> IO a
withMutex (Connection -> Mutex
connMutex Connection
conn) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    current <- IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess)
    case current of
      Just ActiveResponse
ar' | ActiveResponse -> IORef ReadState
arState ActiveResponse
ar' IORef ReadState -> IORef ReadState -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> IORef ReadState
arState ActiveResponse
ar -> do
        ActiveResponse -> IO ()
abandonHandle ActiveResponse
ar
        broken <- Maybe UnusableReason -> Bool
forall a. Maybe a -> Bool
isJust (Maybe UnusableReason -> Bool)
-> IO (Maybe UnusableReason) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Maybe UnusableReason) -> IO (Maybe UnusableReason)
forall a. IORef a -> IO a
readIORef (Connection -> IORef (Maybe UnusableReason)
connBroken Connection
conn)
        unless broken $ void (cancelStatement conn ar)
      Maybe ActiveResponse
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar

-- | Abandon the handle: the session's response gets a new generation, so
-- that reading on through the caller's record throws 'ResultAbandoned'
-- ('checkCurrent'), while the drain, which reads the session's current
-- record, goes on.  The caller holds the lock.
abandonHandle :: ActiveResponse -> IO ()
abandonHandle :: ActiveResponse -> IO ()
abandonHandle ActiveResponse
ar = do
  gen <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Session -> IORef Int
sessGeneration Session
sess) (\Int
g -> (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  writeIORef (sessResponse sess) (Just ar { arGeneration = gen })
  where
    sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar

-- | Send an Attention for the response, unless it has fully arrived (the
-- statement is over, and an Attention would get an acknowledgement of its
-- own) or one was sent already; whether one was sent now.  The caller
-- holds the lock and runs masked.  A failed send breaks the connection; a
-- second interruption, of the send, wins, as in 'bracket'.
cancelStatement :: Connection -> ActiveResponse -> IO Bool
cancelStatement :: Connection -> ActiveResponse -> IO Bool
cancelStatement Connection
conn ActiveResponse
ar = do
  arrived <- ActiveResponse -> IO Bool
fullyArrived ActiveResponse
ar
  attn <- attentionWasSent <$> readIORef (arCancellation ar)
  if arrived || attn
    then pure False
    else do
      ctx <- readIORef (connContext conn)
      sendFramed conn (arSession ar) ctx PTAttention (PacketStatus 0) LB.empty
        `catch` \(SomeException
c :: SomeException) -> case SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
c of
          Just (SomeAsyncException
_ :: SomeAsyncException) -> SomeException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
c
          Maybe SomeAsyncException
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      writeIORef (arCancellation ar) AttentionSent
      pure True

-- | Every session of the connection: the one, or under MARS the open
-- SMP sessions, in the order of their identifiers.
allSessions :: Connection -> IO [Session]
allSessions :: Connection -> IO [Session]
allSessions Connection
conn = case Connection -> Maybe Mars
connMars Connection
conn of
  Maybe Mars
Nothing -> [Session] -> IO [Session]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Connection -> Session
connPrimary Connection
conn]
  Just Mars
mars -> IntMap Session -> [Session]
forall a. IntMap a -> [a]
IntMap.elems (IntMap Session -> [Session])
-> IO (IntMap Session) -> IO [Session]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap Session) -> IO (IntMap Session)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (IntMap Session)
marsOpen Mars
mars)

-- | The session for the next request: the connection's one; under MARS
-- the reserved session ('withReservedSession'), else the idle session
-- with the lowest identifier, else a new one.  The caller holds the lock
-- and registers a response on it, or fails, without being interrupted.
takeSession :: Connection -> IO Session
takeSession :: Connection -> IO Session
takeSession Connection
conn = case Connection -> Maybe Mars
connMars Connection
conn of
  Maybe Mars
Nothing -> Session -> IO Session
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Connection -> Session
connPrimary Connection
conn)
  Just Mars
mars -> IORef (Maybe Session) -> IO (Maybe Session)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars) IO (Maybe Session) -> (Maybe Session -> IO Session) -> IO Session
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just Session
sess -> Session -> IO Session
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Session
sess
    Maybe Session
Nothing -> Connection -> Mars -> IO [Session]
idleSessions Connection
conn Mars
mars IO [Session] -> ([Session] -> IO Session) -> IO Session
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Session
sess : [Session]
_ -> Session -> IO Session
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Session
sess
      [] -> Connection -> Mars -> IO Session
openSession Connection
conn Mars
mars

-- | The open sessions with no response, lowest identifier first, but
-- for a reserved one.
idleSessions :: Connection -> Mars -> IO [Session]
idleSessions :: Connection -> Mars -> IO [Session]
idleSessions Connection
_ Mars
mars = do
  reserved <- (Session -> Int) -> Maybe Session -> Maybe Int
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Session -> Int
sessionId (Maybe Session -> Maybe Int)
-> IO (Maybe Session) -> IO (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Maybe Session) -> IO (Maybe Session)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars)
  open <- IntMap.elems <$> readIORef (marsOpen mars)
  let idle Session
sess
        | Int -> Maybe Int
forall a. a -> Maybe a
Just (Session -> Int
sessionId Session
sess) Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe Int
reserved = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
        | Bool
otherwise = Maybe ActiveResponse -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe ActiveResponse -> Bool)
-> IO (Maybe ActiveResponse) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess)
  filterIO idle open
  where
    filterIO :: (a -> f Bool) -> t a -> f [a]
filterIO a -> f Bool
f = (t [a] -> [a]) -> f (t [a]) -> f [a]
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap t [a] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (f (t [a]) -> f [a]) -> (t a -> f (t [a])) -> t a -> f [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f [a]) -> t a -> f (t [a])
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM (\a
x -> a -> f Bool
f a
x f Bool -> (Bool -> f [a]) -> f [a]
forall a b. f a -> (a -> f b) -> f b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
keep -> [a] -> f [a]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a
x | Bool
keep])

-- | An SMP session's identifier (-1 for the session of a connection
-- without MARS).
sessionId :: Session -> Int
sessionId :: Session -> Int
sessionId = Int -> (SmpSession -> Int) -> Maybe SmpSession -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (-Int
1) (Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16 -> Int) -> (SmpSession -> Word16) -> SmpSession -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmpSession -> Word16
smpSid) (Maybe SmpSession -> Int)
-> (Session -> Maybe SmpSession) -> Session -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Session -> Maybe SmpSession
sessSmp

-- | Open an SMP session: the lowest identifier the server doesn't have
-- open (so at most one above the largest in use, which is all it accepts;
-- an identifier closed by a FIN is free again once the server's FIN has
-- arrived, which takes it out of the registry).  The session is
-- registered before its SYN is sent, so the client can't forget an
-- identifier the server has: a SYN that fails or is interrupted breaks
-- the connection, like any send.  The server doesn't answer a SYN.
openSession :: Connection -> Mars -> IO Session
openSession :: Connection -> Mars -> IO Session
openSession Connection
conn Mars
mars = IO Session -> IO Session
forall a. IO a -> IO a
mask_ (IO Session -> IO Session) -> IO Session -> IO Session
forall a b. (a -> b) -> a -> b
$ do
  links <- TVar (IntMap SmpLink) -> IO (IntMap SmpLink)
forall a. TVar a -> IO a
readTVarIO (Mars -> TVar (IntMap SmpLink)
marsRegistry Mars
mars)
  let sid = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (\Int
i -> Bool -> Bool
not (Int -> IntMap SmpLink -> Bool
forall a. Int -> IntMap a -> Bool
IntMap.member Int
i IntMap SmpLink
links)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0
  when (sid > 0xFFFF) $ throwIO $ EncodeError
    "the connection has 65536 responses open, the most SMP has sessions for"
  link <- SmpLink <$> newTQueueIO <*> newTVarIO smpInitialWindow <*> newTVarIO 0
    <*> newTVarIO (marsWindow mars) <*> newTVarIO 0 <*> newTVarIO False
  smp <- SmpSession (fromIntegral sid) link <$> newIORef 0 <*> newIORef (marsWindow mars)
  sess <- newSession (Just smp)
  atomically $ modifyTVar' (marsRegistry mars) (IntMap.insert sid link)
  modifyIORef' (marsOpen mars) (IntMap.insert sid sess)
  sendSmpControl conn smp SmpSyn
  pure sess

-- | How many idle SMP sessions a connection keeps (SqlClient's number).
maxIdleSessions :: Int
maxIdleSessions :: Int
maxIdleSessions = Int
10

-- | Close idle sessions beyond 'maxIdleSessions', highest identifier
-- first: a FIN, which the server answers with a FIN.  Until that arrives
-- the identifier stays in the registry, neither open nor free for a new
-- session (the server drops the connection on a SYN for it).  Only an
-- idle session is ever closed: a FIN for one whose request the server
-- hasn't finished drops the connection too.
closeIdleSessions :: Connection -> Mars -> IO ()
closeIdleSessions :: Connection -> Mars -> IO ()
closeIdleSessions Connection
conn Mars
mars = do
  idle <- Connection -> Mars -> IO [Session]
idleSessions Connection
conn Mars
mars
  forM_ (take (length idle - maxIdleSessions) (reverse idle)) $ \Session
sess ->
    Maybe SmpSession -> (SmpSession -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Session -> Maybe SmpSession
sessSmp Session
sess) ((SmpSession -> IO ()) -> IO ()) -> (SmpSession -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \SmpSession
smp -> IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (SmpLink -> TVar Bool
linkClosing (SmpSession -> SmpLink
smpLink SmpSession
smp)) Bool
True
      IORef (IntMap Session)
-> (IntMap Session -> IntMap Session) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Mars -> IORef (IntMap Session)
marsOpen Mars
mars) (Int -> IntMap Session -> IntMap Session
forall a. Int -> IntMap a -> IntMap a
IntMap.delete (Session -> Int
sessionId Session
sess))
      Connection -> SmpSession -> SmpFlag -> IO ()
sendSmpControl Connection
conn SmpSession
smp SmpFlag
SmpFin

-- | Start SMP on a connection that negotiated MARS, once it has logged in
-- (and run its session setup) in plain TDS: open session 0, with the
-- given window (at least 4 packets).  Before 'startReader'.  The
-- connection returned is the one to use; the one given is no longer.
enableMars :: Connection -> Int -> IO Connection
enableMars :: Connection -> Int -> IO Connection
enableMars Connection
conn Int
window = do
  leftover <- IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef (Connection -> IORef [ByteString]
primaryBuffer Connection
conn)
  unless (null leftover) $ throwIO $ ProtocolError "unexpected data after the login response"
  marsRegistry <- newTVarIO IntMap.empty
  marsOpen <- newIORef IntMap.empty
  marsReserved <- newIORef Nothing
  marsResetInFlight <- newIORef Nothing
  let mars = Mars { marsWindow :: Word32
marsWindow = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
4 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
0x7FFFFFFF Int
window)), TVar (IntMap SmpLink)
IORef (Maybe ActiveResponse)
IORef (Maybe Session)
IORef (IntMap Session)
marsRegistry :: TVar (IntMap SmpLink)
marsOpen :: IORef (IntMap Session)
marsReserved :: IORef (Maybe Session)
marsResetInFlight :: IORef (Maybe ActiveResponse)
marsRegistry :: TVar (IntMap SmpLink)
marsOpen :: IORef (IntMap Session)
marsReserved :: IORef (Maybe Session)
marsResetInFlight :: IORef (Maybe ActiveResponse)
.. }
      conn' = Connection
conn { connMars = Just mars }
  primary <- openSession conn' mars
  pure conn' { connPrimary = primary }

-- | Run the action with every request it makes sent on one session, taken
-- for it; the lock is held throughout.  For an exchange that spans
-- several requests which the server ties to a session: a bulk load's
-- @INSERT BULK@, its BULK LOAD message, and the Attention that recovers a
-- failed load.  As without MARS, each request inside first finishes the
-- session's previous response.  Without MARS it is 'withConnectionLock'.
withReservedSession :: Connection -> IO a -> IO a
withReservedSession :: forall a. Connection -> IO a -> IO a
withReservedSession Connection
conn IO a
act = Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ case Connection -> Maybe Mars
connMars Connection
conn of
  Maybe Mars
Nothing -> IO a
act
  Just Mars
mars -> IORef (Maybe Session) -> IO (Maybe Session)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars) IO (Maybe Session) -> (Maybe Session -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just Session
_ -> IO a
act
    Maybe Session
Nothing -> IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
      (Connection -> IO ()
checkNotBroken Connection
conn IO () -> IO Session -> IO Session
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Connection -> IO Session
takeSession Connection
conn IO Session -> (Session -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IORef (Maybe Session) -> Maybe Session -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars) (Maybe Session -> IO ())
-> (Session -> Maybe Session) -> Session -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Session -> Maybe Session
forall a. a -> Maybe a
Just)
      (IORef (Maybe Session) -> Maybe Session -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Mars -> IORef (Maybe Session)
marsReserved Mars
mars) Maybe Session
forall a. Maybe a
Nothing)
      IO a
act

-- | Send one message and make its response its session's current one.  If
-- @resettable@, a pending reset is consumed and the message carries the
-- RESETCONNECTION bit.
--
-- Returns normally if and only if the message was sent and its response
-- registered; if the send fails or is interrupted, the connection is broken
-- ('sendBytes').  A failure before anything is sent (an encoding error)
-- leaves the connection usable, and a pending reset pending.
sendRequest :: Connection -> Maybe Text -> PacketType -> Bool -> (TdsContext -> PacketWriter -> IO ()) -> IO ActiveResponse
sendRequest :: Connection
-> Maybe Text
-> PacketType
-> Bool
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
sendRequest Connection
conn Maybe Text
stmt PacketType
ptype Bool
resettable TdsContext -> PacketWriter -> IO ()
writeBody = IO ActiveResponse -> IO ActiveResponse
forall a. IO a -> IO a
mask_ (IO ActiveResponse -> IO ActiveResponse)
-> IO ActiveResponse -> IO ActiveResponse
forall a b. (a -> b) -> a -> b
$ do
  ctx <- IORef TdsContext -> IO TdsContext
forall a. IORef a -> IO a
readIORef (Connection -> IORef TdsContext
connContext Connection
conn)
  sess <- takeSession conn
  resetting <- if resettable then atomicModifyIORef' (connResetPending conn) (False,) else pure False
  let status = if Bool
resetting then PacketStatus
statusResetConnection else Word8 -> PacketStatus
PacketStatus Word8
0
  ar0 <- newResponse sess stmt resetting
  sent <- newIORef False
  result <- sendPacketsTracked conn sess ctx ptype status (writeIORef sent True) (writeBody ctx)
  case result of
    Right () -> Connection -> ActiveResponse -> IO ActiveResponse
registerResponse Connection
conn ActiveResponse
ar0
    Left SomeException
err -> do
      anyPacket <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
sent
      if not anyPacket && ptype /= PTBulkLoad then do
        when resetting (writeIORef (connResetPending conn) True)
        throwIO err
      else do
        ar <- registerResponse conn ar0
        sendFramed conn sess ctx PTAttention (PacketStatus 0) LB.empty
        writeIORef (arCancellation ar) AttentionSent
        finishSession conn sess
        throwIO err

-- | The response to a message about to be sent on the session, before its
-- generation: 'registerResponse' gives it that once the message is sent.
newResponse :: Session -> Maybe Text -> Bool -> IO ActiveResponse
newResponse :: Session -> Maybe Text -> Bool -> IO ActiveResponse
newResponse Session
sess Maybe Text
stmt Bool
resetting = do
  arEomBase <- IO Int -> (SmpSession -> IO Int) -> Maybe SmpSession -> IO Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0) (TVar Int -> IO Int
forall a. TVar a -> IO a
readTVarIO (TVar Int -> IO Int)
-> (SmpSession -> TVar Int) -> SmpSession -> IO Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmpLink -> TVar Int
linkEoms (SmpLink -> TVar Int)
-> (SmpSession -> SmpLink) -> SmpSession -> TVar Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmpSession -> SmpLink
smpLink) (Session -> Maybe SmpSession
sessSmp Session
sess)
  arSawResetAck <- newIORef False
  arCancellation <- newIORef Running
  arInConsumer <- newIORef False
  arReplay <- newIORef []
  arState <- newIORef (ReadState False (AtToken B.empty))
  arResultSet <- newIORef Nothing
  arSawFinalDone <- newIORef False
  pure ActiveResponse
    { arGeneration = 0, arSession = sess, arStatement = stmt, arCarriedReset = resetting, .. }

-- | Make the response to the message just sent its session's current one,
-- as a new generation of the session.  Under MARS, one that carried a
-- reset holds back the requests after it ('awaitReset').
registerResponse :: Connection -> ActiveResponse -> IO ActiveResponse
registerResponse :: Connection -> ActiveResponse -> IO ActiveResponse
registerResponse Connection
conn ActiveResponse
ar0 = do
  let sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar0
  gen <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Session -> IORef Int
sessGeneration Session
sess) (\Int
g -> (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  let ar = ActiveResponse
ar0 { arGeneration = gen }
  writeIORef (sessResponse sess) (Just ar)
  when (arCarriedReset ar) $ forM_ (connMars conn) $ \Mars
mars ->
    IORef (Maybe ActiveResponse) -> Maybe ActiveResponse -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight Mars
mars) (ActiveResponse -> Maybe ActiveResponse
forall a. a -> Maybe a
Just ActiveResponse
ar)
  pure ar

-- | Under MARS, a request sent on another session right behind one that
-- carries RESETCONNECTION races it: the server may run either first, and
-- refuses what it runs before the reset.  The reset's acknowledgement is
-- among the ENVCHANGE tokens that begin the carrying request's response,
-- and nothing sent once those have arrived is refused.  So before a
-- request starts, the response of a reset still in flight is read up to
-- its first token that is not an ENVCHANGE (or its end).
--
-- The tokens are read by the ordinary reader, which applies each once,
-- from wherever the response's consumer has got to, and kept for the
-- consumer, which is given them first ('arReplay').  The wait comes
-- before anything is sent, so an interruption is harmless: it doesn't
-- cancel the response, and the tokens read so far are kept.
awaitReset :: Connection -> Mars -> IO ()
awaitReset :: Connection -> Mars -> IO ()
awaitReset Connection
conn Mars
mars = IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight Mars
mars) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Maybe ActiveResponse
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Just ActiveResponse
ar -> do
    current <- IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse (ActiveResponse -> Session
arSession ActiveResponse
ar))
    case current of
      Just ActiveResponse
ar' | ActiveResponse -> Int
arGeneration ActiveResponse
ar' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> Int
arGeneration ActiveResponse
ar -> do
        -- Masked from consuming a token to keeping it; the wait for its
        -- packet stays interruptible.
        step <- IO (ResponseStep (ZonkAny 7)) -> IO (ResponseStep (ZonkAny 7))
forall a. IO a -> IO a
mask_ (IO (ResponseStep (ZonkAny 7)) -> IO (ResponseStep (ZonkAny 7)))
-> IO (ResponseStep (ZonkAny 7)) -> IO (ResponseStep (ZonkAny 7))
forall a b. (a -> b) -> a -> b
$ do
          step <- Connection
-> ActiveResponse
-> Bool
-> RowRead (ZonkAny 7)
-> IO (ResponseStep (ZonkAny 7))
forall r.
Connection
-> ActiveResponse -> Bool -> RowRead r -> IO (ResponseStep r)
stepCore Connection
conn ActiveResponse
ar Bool
False RowRead (ZonkAny 7)
forall a. RowRead a
StopAtRow
          modifyIORef' (arReplay ar) (++ replayed step)
          settleReset conn ar step
          pure step
        runCallbacks conn step
        awaitReset conn mars
      -- Canceled, or drained: whoever did that has read it.
      Maybe ActiveResponse
_ -> IORef (Maybe ActiveResponse) -> Maybe ActiveResponse -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight Mars
mars) Maybe ActiveResponse
forall a. Maybe a
Nothing
  where
    replayed :: ResponseStep r -> [Token]
replayed = \case
      StepToken Token
tok -> [Token
tok]
      StepRows RowBatch r
_ -> []
      ResponseStep r
StepEnd -> []

-- | After a step of a response that carried a reset: once it is past its
-- leading ENVCHANGE tokens, requests may start again ('awaitReset').
settleReset :: Connection -> ActiveResponse -> ResponseStep r -> IO ()
settleReset :: forall r. Connection -> ActiveResponse -> ResponseStep r -> IO ()
settleReset Connection
conn ActiveResponse
ar ResponseStep r
step = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ActiveResponse -> Bool
arCarriedReset ActiveResponse
ar) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ case ResponseStep r
step of
  StepToken (TEnvChange EnvChange
_) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  ResponseStep r
_ -> Maybe Mars -> (Mars -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Connection -> Maybe Mars
connMars Connection
conn) ((Mars -> IO ()) -> IO ()) -> (Mars -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Mars
mars -> IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight Mars
mars) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just ActiveResponse
inFlight | ActiveResponse -> IORef ReadState
arState ActiveResponse
inFlight IORef ReadState -> IORef ReadState -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> IORef ReadState
arState ActiveResponse
ar -> IORef (Maybe ActiveResponse) -> Maybe ActiveResponse -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Mars -> IORef (Maybe ActiveResponse)
marsResetInFlight Mars
mars) Maybe ActiveResponse
forall a. Maybe a
Nothing
    Maybe ActiveResponse
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | 'startRequest' for a message produced packet by packet as it is sent
-- ('sendPackets'): the producer gets a 'PacketWriter', and the message ends
-- when it returns.  A pending 'requestReset' is
-- left pending (the message isn't a SQL batch or RPC).
--
-- If the producer fails (synchronously) with the message unfinished, the
-- request is canceled: an Attention is sent, and the server's
-- acknowledgement read (as after an interrupted read, 'nextToken'), so
-- the connection stays usable; the producer's exception is then
-- rethrown.  A failed or interrupted send breaks the connection, as any
-- send does.
startStreamingRequest :: Connection -> Maybe Text -> PacketType -> (PacketWriter -> IO ()) -> IO ActiveResponse
startStreamingRequest :: Connection
-> Maybe Text
-> PacketType
-> (PacketWriter -> IO ())
-> IO ActiveResponse
startStreamingRequest Connection
conn Maybe Text
stmt PacketType
ptype PacketWriter -> IO ()
produce = Connection
-> Maybe Text
-> PacketType
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
startRequestFor Connection
conn Maybe Text
stmt PacketType
ptype (\TdsContext
_ -> PacketWriter -> IO ()
produce)

-- | Reset the session and run the session setup again, in one SQLBatch that
-- carries the RESETCONNECTION bit, read to the end.  SET statements inside
-- an RPC (such as @sp_executesql@) revert when it returns, and prefixing
-- them to the caller's batch would break batches that must start with
-- CREATE PROCEDURE and the like, so this costs a round trip of its own.
--
-- If reading the response fails or is interrupted, it is unknown how much
-- of the setup ran, so the connection is marked broken; likewise if the
-- setup fails on the server (as in 'Hmssql.Connection.connect', a session
-- without its setup is not usable): with an ERROR token ('SqlServerError'),
-- or with a DONE flagged as failed and no ERROR token ('SevereError').
--
-- Unlike @sp_reset_connection@, RESETCONNECTION keeps the transaction
-- isolation level (SQL Server 2025; Azure SQL Database puts READ COMMITTED
-- back); nothing here restores it.
resetWithSessionSetup :: Connection -> IO ()
resetWithSessionSetup :: Connection -> IO ()
resetWithSessionSetup Connection
conn = ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
  let sql :: Text
sql = Text -> [Text] -> Text
T.intercalate Text
"\n" (Connection -> [Text]
connSessionSetup Connection
conn)
  ar <- Connection
-> Maybe Text
-> PacketType
-> Bool
-> (TdsContext -> PacketWriter -> IO ())
-> IO ActiveResponse
sendRequest Connection
conn (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
sql) PacketType
PTSqlBatch Bool
True (\TdsContext
ctx PacketWriter
writer -> do
    ([Char] -> IO (ZonkAny 8)) -> Maybe [Char] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (EncodeError -> IO (ZonkAny 8)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (EncodeError -> IO (ZonkAny 8))
-> ([Char] -> EncodeError) -> [Char] -> IO (ZonkAny 8)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> EncodeError
EncodeError) (TdsContext -> Text -> Maybe [Char]
sqlBatchProblem TdsContext
ctx Text
sql)
    PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
writer (TdsContext -> Text -> Builder
sqlBatchBuilder TdsContext
ctx Text
sql))
  let go Maybe ServerMessage
err Maybe Done
done = Connection -> ActiveResponse -> IO (Maybe Token)
nextToken Connection
conn ActiveResponse
ar IO (Maybe Token)
-> (Maybe Token -> IO (Maybe ServerMessage, Maybe Done))
-> IO (Maybe ServerMessage, Maybe Done)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe Token
Nothing -> (Maybe ServerMessage, Maybe Done)
-> IO (Maybe ServerMessage, Maybe Done)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ServerMessage
err, Maybe Done
done)
        Just (TError ServerMessage
msg) -> Maybe ServerMessage
-> Maybe Done -> IO (Maybe ServerMessage, Maybe Done)
go (Maybe ServerMessage
err Maybe ServerMessage -> Maybe ServerMessage -> Maybe ServerMessage
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ServerMessage -> Maybe ServerMessage
forall a. a -> Maybe a
Just ServerMessage
msg) Maybe Done
done
        Just Token
tok -> Maybe ServerMessage
-> Maybe Done -> IO (Maybe ServerMessage, Maybe Done)
go Maybe ServerMessage
err (Maybe Done
done Maybe Done -> Maybe Done -> Maybe Done
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Token -> Maybe Done
failedDoneOf Token
tok)
  failed <- restore (go Nothing Nothing) `catch` \(SomeException
e :: SomeException) -> do
    Connection -> UnusableReason -> IO ()
markBroken Connection
conn ((SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
ResetFailed SomeException
e)
    SomeException -> IO (Maybe ServerMessage, Maybe Done)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
  -- A failed DONE with no ERROR token is a failure too: it is how the
  -- server refuses a reset (under MARS, with another request active).
  let failure = case (Maybe ServerMessage, Maybe Done)
failed of
        (Just ServerMessage
msg, Maybe Done
_) -> SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (SqlServerError -> SomeException
forall e. Exception e => e -> SomeException
toException (ServerMessage -> Maybe Text -> SqlServerError
SqlServerError ServerMessage
msg (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
sql)))
        (Maybe ServerMessage
Nothing, Just Done
d) -> SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (SevereError -> SomeException
forall e. Exception e => e -> SomeException
toException (Done -> Maybe Text -> SevereError
SevereError Done
d (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
sql)))
        (Maybe ServerMessage
Nothing, Maybe Done
Nothing) -> Maybe SomeException
forall a. Maybe a
Nothing
  forM_ failure $ \SomeException
e -> do
    Connection -> UnusableReason -> IO ()
markBroken Connection
conn (SomeException -> UnusableReason
resetFailure SomeException
e)
    SomeException -> IO (ZonkAny 9)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e

-- | Reset the session before the next SQLBatch, RPC or transaction manager
-- request: the session returns to its state right after login (MS-TDS
-- 2.2.3.1.2, RESETCONNECTION), and the session setup runs again.  With no
-- session setup, the next request itself carries the RESETCONNECTION bit,
-- so the reset costs nothing extra; otherwise it is preceded by one batch
-- that carries the bit and the session setup.
--
-- Takes the connection lock, so a reset requested by another thread waits
-- for whatever holds the connection (a request being read, or a whole
-- 'Hmssql.Transaction.withTransaction') and applies to the request after it.
requestReset :: Connection -> IO ()
requestReset :: Connection -> IO ()
requestReset Connection
conn = Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
  IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
atomicWriteIORef (Connection -> IORef Bool
connResetPending Connection
conn) Bool
True

-- | The rows at the head of the response, decoded: what one 'nextStep'
-- returns when it finds rows.
data RowBatch r = RowBatch
  { forall r. RowBatch r -> [r]
batchRows :: ![r]
    -- ^ In order.
  , forall r. RowBatch r -> Maybe DecodeError
batchFailed :: !(Maybe DecodeError)
    -- ^ The row after the last one didn't decode; it has been consumed,
    -- and the response goes on after it.
  }

-- | Where the reading of 'stepCore' stopped: at a step to return, or at a consumer to
-- run before reading on.  A consumer is user code, so the reading hands it out
-- rather than running it inside the handler that breaks the connection when the
-- bytes can't be read.  'StepConsumer' holds the consumer and the reading that
-- goes on with its result, as 'P.RunConsumer' does for a recipe.
data StepProgress r
  = StepReady (ResponseStep r)
  | forall x. StepConsumer (IO x) (x -> IO (StepProgress r))

-- | What 'nextStep' found next in the response.
data ResponseStep r
  = StepRows !(RowBatch r)
  | StepToken !Token
    -- ^ A token other than a row.
  | StepEnd
    -- ^ The response is complete.

-- | Read a batch of rows, a non-row token, or the end of the response.
-- Complete rows are returned before waiting for another packet. A partial
-- token remains as a continuation in the response, without re-reading bytes.
nextStep :: Connection -> ActiveResponse -> RowRead r -> IO (ResponseStep r)
nextStep :: forall r.
Connection -> ActiveResponse -> RowRead r -> IO (ResponseStep r)
nextStep Connection
conn ActiveResponse
ar RowRead r
decodeRow = Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepResponse Bool
False Connection
conn ActiveResponse
ar Bool
False RowRead r
decodeRow

-- 'nextStep' with each batch's rows in reverse order, as they were accumulated,
-- for 'Hmssql.Query.Internal.readFirstResultSet', which collects every batch and
-- reverses once at the end.  No other caller may use it: a 'RowBatch' is in order
-- everywhere else.
nextStepReversed :: Connection -> ActiveResponse -> RowRead r -> IO (ResponseStep r)
nextStepReversed :: forall r.
Connection -> ActiveResponse -> RowRead r -> IO (ResponseStep r)
nextStepReversed Connection
conn ActiveResponse
ar RowRead r
decodeRow = Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepResponse Bool
True Connection
conn ActiveResponse
ar Bool
False RowRead r
decodeRow

-- | Read one raw token. An interrupted receive sends Attention eagerly
-- unless the response has already arrived, and its continuation is retained
-- for the drain. A malformed token breaks the connection.
nextToken :: Connection -> ActiveResponse -> IO (Maybe Token)
nextToken :: Connection -> ActiveResponse -> IO (Maybe Token)
nextToken Connection
conn ActiveResponse
ar = Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead Row
-> IO (ResponseStep Row)
forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepResponse Bool
False Connection
conn ActiveResponse
ar Bool
True RowRead Row
wholeRows IO (ResponseStep Row)
-> (ResponseStep Row -> IO (Maybe Token)) -> IO (Maybe Token)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  StepRows (RowBatch [Row
row] Maybe DecodeError
_) -> Maybe Token -> IO (Maybe Token)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Token -> Maybe Token
forall a. a -> Maybe a
Just (Row -> Token
TRow Row
row))
  StepRows RowBatch Row
_ -> ProtocolError -> IO (Maybe Token)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (ProtocolError -> IO (Maybe Token))
-> ProtocolError -> IO (Maybe Token)
forall a b. (a -> b) -> a -> b
$ [Char] -> ProtocolError
ProtocolError [Char]
"nextToken: a batch of one row without the row"  -- can't happen: at most one row, decoded with Right
  StepToken Token
tok -> Maybe Token -> IO (Maybe Token)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Token -> Maybe Token
forall a. a -> Maybe a
Just Token
tok)
  ResponseStep Row
StepEnd -> Maybe Token -> IO (Maybe Token)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Token
forall a. Maybe a
Nothing

-- | Throw 'ResultAbandoned' unless this response is still its session's
-- current one: not superseded by a later request (without MARS), canceled
-- by an interruption or by 'cancelActiveResponse', drained, or closed
-- with the connection (which throws its own reason);
-- 'ConnectionReleased' if the handle was released ('checkLease').  What 'nextStep' checks before reading; a
-- stream that holds rows of a batch checks it before yielding each, so
-- that a superseded stream throws at once, not after the rows it had.
checkCurrent :: Connection -> ActiveResponse -> IO ()
checkCurrent :: Connection -> ActiveResponse -> IO ()
checkCurrent Connection
conn ActiveResponse
ar = do
  IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef Bool
arInConsumer ActiveResponse
ar) IO Bool -> (Bool -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
active -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
active (ReentrantCallback -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ReentrantCallback
ReentrantCallback)
  Connection -> IO ()
checkLease Connection
conn
  Connection -> IO ()
checkNotBroken Connection
conn
  IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse (ActiveResponse -> Session
arSession ActiveResponse
ar)) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just ActiveResponse
ar' | ActiveResponse -> Int
arGeneration ActiveResponse
ar' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> Int
arGeneration ActiveResponse
ar -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Maybe ActiveResponse
_ -> Connection -> IO ()
forall a. Connection -> IO a
abandoned Connection
conn

-- | Throw 'ResultAbandoned' for a response that is no longer the current
-- one, unless the connection was closed or broke meanwhile: 'close' from
-- another thread (which doesn't wait for the connection lock) may record
-- its reason and clear the response ('releaseTransport') just after the
-- caller's 'checkNotBroken'.  Every use after 'close' throws its reason.
abandoned :: Connection -> IO a
abandoned :: forall a. Connection -> IO a
abandoned Connection
conn = Connection -> IO ()
checkNotBroken Connection
conn IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ResultAbandoned -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ResultAbandoned
ResultAbandoned

-- | 'nextStep', optionally restricted to one row for 'nextToken'.
stepResponse :: Bool -> Connection -> ActiveResponse -> Bool -> RowRead r -> IO (ResponseStep r)
stepResponse :: forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepResponse Bool
reversed Connection
conn ActiveResponse
ar Bool
oneRow RowRead r
decodeRow = Connection -> IO (ResponseStep r) -> IO (ResponseStep r)
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO (ResponseStep r) -> IO (ResponseStep r))
-> IO (ResponseStep r) -> IO (ResponseStep r)
forall a b. (a -> b) -> a -> b
$ do
  IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef Bool
arInConsumer ActiveResponse
ar) IO Bool -> (Bool -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
active -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
active (ReentrantCallback -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ReentrantCallback
ReentrantCallback)
  Connection -> IO ()
checkNotBroken Connection
conn
  current <- IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess)
  case current of
    Just ActiveResponse
ar' | ActiveResponse -> Int
arGeneration ActiveResponse
ar' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> Int
arGeneration ActiveResponse
ar ->
      -- Only a response that carried a reset has tokens read ahead
      -- ('awaitReset').
      if ActiveResponse -> Bool
arCarriedReset ActiveResponse
ar then IORef [Token] -> IO [Token]
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef [Token]
arReplay ActiveResponse
ar) IO [Token]
-> ([Token] -> IO (ResponseStep r)) -> IO (ResponseStep r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Token] -> IO (ResponseStep r)
replayOr else IO (ResponseStep r)
step
    Maybe ActiveResponse
_ -> Connection -> IO (ResponseStep r)
forall a. Connection -> IO a
abandoned Connection
conn
  where
    sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar

    -- The token is fully consumed before the callbacks run, so a failing
    -- callback doesn't leave the byte stream in an unknown position; an
    -- interruption inside a callback still cancels the statement.
    step :: IO (ResponseStep r)
step = (do
        r <- Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepCoreWithOrder Bool
reversed Connection
conn ActiveResponse
ar Bool
oneRow RowRead r
decodeRow
        settleReset conn ar r
        runCallbacks conn r
        pure r)
      IO (ResponseStep r)
-> (SomeAsyncException -> IO (ResponseStep r))
-> IO (ResponseStep r)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeAsyncException
e :: SomeAsyncException) -> do
        IO ()
cancelEagerly
        SomeAsyncException -> IO (ResponseStep r)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeAsyncException
e

    -- A token 'awaitReset' read: its effects are applied and its
    -- callbacks have run.
    replayOr :: [Token] -> IO (ResponseStep r)
replayOr = \case
      [] -> IO (ResponseStep r)
step
      Token
tok : [Token]
rest -> IORef [Token] -> [Token] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef [Token]
arReplay ActiveResponse
ar) [Token]
rest IO () -> IO (ResponseStep r) -> IO (ResponseStep r)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ResponseStep r -> IO (ResponseStep r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Token -> ResponseStep r
forall r. Token -> ResponseStep r
StepToken Token
tok)

    -- Runs masked (in a handler).  Once the statement is canceled, the
    -- handle is abandoned: read on through it, the response would end at
    -- the acknowledgement, silently missing the rest of the result.  A
    -- response that had fully arrived is left readable: nothing was sent,
    -- and reading on gives the rest of the result (a batch committed as
    -- consumed before the interruption landed is lost, as one interrupted
    -- on its way back to the caller is; design.md, "Interruption and
    -- cancellation").
    cancelEagerly :: IO ()
cancelEagerly = do
      sent <- Connection -> ActiveResponse -> IO Bool
cancelStatement Connection
conn ActiveResponse
ar
      when sent $ abandonHandle ar

-- | Whether the response's last packet has been received: read by its
-- consumer, or, on an SMP session, queued for it by the reader thread.
-- The statement is then over on the server, and the rest of the response
-- can be read without canceling anything.
fullyArrived :: ActiveResponse -> IO Bool
fullyArrived :: ActiveResponse -> IO Bool
fullyArrived ActiveResponse
ar = do
  eom <- ReadState -> Bool
readFinalPacket (ReadState -> Bool) -> IO ReadState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef ReadState -> IO ReadState
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar)
  case sessSmp (arSession ar) of
    Just SmpSession
smp | Bool -> Bool
not Bool
eom -> (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> ActiveResponse -> Int
arEomBase ActiveResponse
ar) (Int -> Bool) -> IO Int -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar Int -> IO Int
forall a. TVar a -> IO a
readTVarIO (SmpLink -> TVar Int
linkEoms (SmpSession -> SmpLink
smpLink SmpSession
smp))
    Maybe SmpSession
_ -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
eom

-- | Run the connection's callbacks for a step: with the lock held, and
-- the connection refusing to be used meanwhile.
runCallbacks :: Connection -> ResponseStep r -> IO ()
runCallbacks :: forall r. Connection -> ResponseStep r -> IO ()
runCallbacks Connection
conn = \case
  StepToken (TEnvChange EnvChange
ec) -> IO () -> IO ()
forall a. IO a -> IO a
inCallback (Connection -> EnvChange -> IO ()
connOnEnvChange Connection
conn EnvChange
ec)
  StepToken (TInfo ServerMessage
msg) -> IO () -> IO ()
forall a. IO a -> IO a
inCallback (Connection -> ServerMessage -> IO ()
connOnInfo Connection
conn ServerMessage
msg)
  ResponseStep r
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    inCallback :: IO c -> IO c
inCallback =
      IO () -> IO () -> IO c -> IO c
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ (IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
connInCallback Connection
conn) Bool
True) (IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
connInCallback Connection
conn) Bool
False)

-- | The reading of 'stepResponse': the next step of the response, from
-- the bytes pending or the session's next packets.  The caller holds the
-- lock and has checked that the response is its session's current one.
-- Completed rows are returned before a receive would block. State is committed
-- at the batch boundary or before a suspension, never for each completed row.
--
-- A batch is the rows that complete with the bytes in hand.  When a row runs out
-- of bytes with others already decoded, the batch ends and the suspended row is
-- kept, so complete rows never wait behind a slow server, and the next step goes
-- on with that row rather than reading it again.
--
-- The whole step runs masked, and only the wait for a packet is interruptible.  An
-- interruption there finds the suspension already committed, so nothing is lost.
-- A consumer is user code, so it alone runs in the caller's masking state.
stepCore :: Connection -> ActiveResponse -> Bool -> RowRead r -> IO (ResponseStep r)
stepCore :: forall r.
Connection
-> ActiveResponse -> Bool -> RowRead r -> IO (ResponseStep r)
stepCore = Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepCoreWithOrder Bool
False

stepCoreWithOrder :: Bool -> Connection -> ActiveResponse -> Bool -> RowRead r -> IO (ResponseStep r)
stepCoreWithOrder :: forall r.
Bool
-> Connection
-> ActiveResponse
-> Bool
-> RowRead r
-> IO (ResponseStep r)
stepCoreWithOrder Bool
reversed Connection
conn ActiveResponse
ar Bool
requestedOne RowRead r
reader = ((forall a. IO a -> IO a) -> IO (ResponseStep r))
-> IO (ResponseStep r)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (ResponseStep r))
 -> IO (ResponseStep r))
-> ((forall a. IO a -> IO a) -> IO (ResponseStep r))
-> IO (ResponseStep r)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
  let sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar
      ahead :: Int
ahead = Connection -> Int
connValueAllocAhead Connection
conn
      oneRow :: Bool
oneRow = Bool
requestedOne
      rows :: ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
state [r]
acc Maybe DecodeError
failed = do
        IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
state
        StepProgress r -> IO (StepProgress r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResponseStep r -> StepProgress r
forall r. ResponseStep r -> StepProgress r
StepReady (RowBatch r -> ResponseStep r
forall r. RowBatch r -> ResponseStep r
StepRows ([r] -> Maybe DecodeError -> RowBatch r
forall r. [r] -> Maybe DecodeError -> RowBatch r
RowBatch (if Bool
reversed then [r]
acc else [r] -> [r]
forall a. [a] -> [a]
reverse [r]
acc) Maybe DecodeError
failed)))
      loop :: [r] -> IO (StepProgress r)
loop [r]
acc = IORef ReadState -> IO ReadState
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) IO ReadState
-> (ReadState -> IO (StepProgress r)) -> IO (StepProgress r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [r] -> ReadState -> IO (StepProgress r)
at [r]
acc
      at :: [r] -> ReadState -> IO (StepProgress r)
at [r]
acc ReadState
state = case ReadState -> ReadPosition
readPosition ReadState
state of
        AtToken ByteString
bytes -> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
atBytes [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
bytes Int
0
        Reading Stream Await IO ReadEvent
suspended -> [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [r]
acc ReadState
state Stream Await IO ReadEvent
suspended
        Parked ByteString
bytes Parser ()
remainder -> [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [r]
acc ReadState
state (Stream Await IO ReadEvent -> IO (StepProgress r))
-> Stream Await IO ReadEvent -> IO (StepProgress r)
forall a b. (a -> b) -> a -> b
$
          (Got () -> ReadEvent)
-> Stream Await IO (Got ()) -> Stream Await IO ReadEvent
forall a b. (a -> b) -> Stream Await IO a -> Stream Await IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(P.Got () ByteString
rest) -> ByteString -> ReadEvent
SkippedRow ByteString
rest) (Parser () -> ByteString -> Stream Await IO (Got ())
forall a. Parser a -> ByteString -> Stream Await IO (Got a)
P.parse Parser ()
remainder ByteString
bytes)
      -- Keep ordinary progress as arguments: the payload and a position in it,
      -- which a completed row returns ('P.runParserAt').  The unread bytes are
      -- built only when a batch or a suspended computation is committed to the
      -- response, or a token other than a row is read; built after every row,
      -- they were 32 bytes of a typed (Int32, Int64) row's 366.
      atBytes :: [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
atBytes [r]
acc Bool
final ByteString
payload !Int
pos
        | Bool
oneRow Bool -> Bool -> Bool
&& Bool -> Bool
not ([r] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [r]
acc) = ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
state [r]
acc Maybe DecodeError
forall a. Maybe a
Nothing
        | ByteString -> Int -> Bool
isRowTokenAt ByteString
payload Int
pos = [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
readRow [r]
acc Bool
final ByteString
payload Int
pos
        | Bool -> Bool
not ([r] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [r]
acc) = ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
state [r]
acc Maybe DecodeError
forall a. Maybe a
Nothing
        | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= ByteString -> Int
B.length ByteString
payload = if Bool
final then IO (StepProgress r)
messageEnd else do
            (body, end) <- Connection -> Session -> IO (ByteString, Bool)
readResponsePacket Connection
conn Session
sess
            atBytes acc end body 0
        | Bool
otherwise = [r] -> ReadState -> ByteString -> IO (StepProgress r)
startToken [r]
acc ReadState
state (Int -> ByteString -> ByteString
B.drop Int
pos ByteString
payload)
        where state :: ReadState
state = Bool -> ReadPosition -> ReadState
ReadState Bool
final (ByteString -> ReadPosition
AtToken (Int -> ByteString -> ByteString
B.drop Int
pos ByteString
payload))
      changed :: IO a
changed = ProtocolError -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ([Char] -> ProtocolError
ProtocolError [Char]
"row decoder changed inside a row")
      -- The fast path is 'P.runParserAt', with no Stream constructor and no handler
      -- per row (the one handler is around the whole step, below).  This is
      -- selected once per step, not per row.
      readRow :: [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
readRow = case RowRead r
reader of
        ReadRows Parser r
recipe IORef (Maybe (Either ReadFailure r))
inbox ->
          \[r]
acc Bool
final ByteString
payload Int
pos -> case Parser r
-> ByteString
-> Int
-> Either (Stream Await IO (Got (Either ReadFailure r))) (At r)
forall a.
Parser a
-> ByteString
-> Int
-> Either (Stream Await IO (Got (Either ReadFailure a))) (At a)
P.runParserAt Parser r
recipe ByteString
payload Int
pos of
            Right (P.At r
result Int
end) -> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
atBytes (r
result r -> [r] -> [r]
forall a. a -> [a] -> [a]
: [r]
acc) Bool
final ByteString
payload Int
end
            Right (P.FailedAt ReadFailure
failure Int
end) -> [r]
-> Bool -> ByteString -> Int -> ReadFailure -> IO (StepProgress r)
failedRow [r]
acc Bool
final ByteString
payload Int
end ReadFailure
failure
            Left Stream Await IO (Got (Either ReadFailure r))
suspended -> [r]
-> ReadState
-> IORef (Maybe (Either ReadFailure r))
-> Stream Await IO (Got (Either ReadFailure r))
-> IO (StepProgress r)
driveRow [r]
acc (Bool -> ReadPosition -> ReadState
ReadState Bool
final (ByteString -> ReadPosition
AtToken (Int -> ByteString -> ByteString
B.drop Int
pos ByteString
payload))) IORef (Maybe (Either ReadFailure r))
inbox Stream Await IO (Got (Either ReadFailure r))
suspended
        WholeRows Row -> r
_ -> ([r] -> ReadState -> ByteString -> IO (StepProgress r))
-> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
forall {t} {t}.
(t -> ReadState -> ByteString -> t)
-> t -> Bool -> ByteString -> Int -> t
withPosition [r] -> ReadState -> ByteString -> IO (StepProgress r)
startToken
        DiscardRows r
_ -> ([r] -> ReadState -> ByteString -> IO (StepProgress r))
-> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
forall {t} {t}.
(t -> ReadState -> ByteString -> t)
-> t -> Bool -> ByteString -> Int -> t
withPosition [r] -> ReadState -> ByteString -> IO (StepProgress r)
skipCurrent
        DrainRows r
_ -> ([r] -> ReadState -> ByteString -> IO (StepProgress r))
-> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
forall {t} {t}.
(t -> ReadState -> ByteString -> t)
-> t -> Bool -> ByteString -> Int -> t
withPosition [r] -> ReadState -> ByteString -> IO (StepProgress r)
skipCurrent
        -- 'awaitReset' stops here with nothing consumed.  Reading a row only to
        -- replay it meant writing it back to TDS and parsing it again.
        RowRead r
StopAtRow -> ([r] -> ReadState -> ByteString -> IO (StepProgress r))
-> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
forall {t} {t}.
(t -> ReadState -> ByteString -> t)
-> t -> Bool -> ByteString -> Int -> t
withPosition (\[r]
acc ReadState
state ByteString
_ -> ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
state [r]
acc Maybe DecodeError
forall a. Maybe a
Nothing)
      withPosition :: (t -> ReadState -> ByteString -> t)
-> t -> Bool -> ByteString -> Int -> t
withPosition t -> ReadState -> ByteString -> t
next t
acc Bool
final ByteString
payload Int
pos =
        let bytes :: ByteString
bytes = Int -> ByteString -> ByteString
B.drop Int
pos ByteString
payload in t -> ReadState -> ByteString -> t
next t
acc (Bool -> ReadPosition -> ReadState
ReadState Bool
final (ByteString -> ReadPosition
AtToken ByteString
bytes)) ByteString
bytes
      -- A row that didn't decode, read to its end.
      failedRow :: [r]
-> Bool -> ByteString -> Int -> ReadFailure -> IO (StepProgress r)
failedRow [r]
acc Bool
final ByteString
payload Int
end = \case
        DecodeFailed DecodeError
err -> ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows (Bool -> ReadPosition -> ReadState
ReadState Bool
final (ByteString -> ReadPosition
AtToken (Int -> ByteString -> ByteString
B.drop Int
end ByteString
payload))) [r]
acc (DecodeError -> Maybe DecodeError
forall a. a -> Maybe a
Just DecodeError
err)
        -- Only a drain discards, and a drain doesn't look at rows.
        ReadFailure
ReadDiscarded -> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
atBytes [r]
acc Bool
final ByteString
payload Int
end
      skipCurrent :: [r] -> ReadState -> ByteString -> IO (StepProgress r)
skipCurrent [r]
acc ReadState
state ByteString
bytes = do
        info <- IORef (Maybe ResultSetInfo) -> IO (Maybe ResultSetInfo)
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef (Maybe ResultSetInfo)
arResultSet ActiveResponse
ar) IO (Maybe ResultSetInfo)
-> (Maybe ResultSetInfo -> IO ResultSetInfo) -> IO ResultSetInfo
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO ResultSetInfo
-> (ResultSetInfo -> IO ResultSetInfo)
-> Maybe ResultSetInfo
-> IO ResultSetInfo
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
          (ProtocolError -> IO ResultSetInfo
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ([Char] -> ProtocolError
ProtocolError [Char]
"ROW without COLMETADATA")) ResultSetInfo -> IO ResultSetInfo
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        resume acc state $ fmap (\(P.Got () ByteString
rest) -> ByteString -> ReadEvent
SkippedRow ByteString
rest) (P.parse (skipRow info) bytes)
      skipRow :: ResultSetInfo -> Parser ()
skipRow ResultSetInfo
info = do
        bitmap <- Int -> Parser ByteString
V.rowBitmap (ResultSetInfo -> Int
rsColumnCount ResultSetInfo
info)
        V.skipCells bitmap (rsCells info)
      startToken :: [r] -> ReadState -> ByteString -> IO (StepProgress r)
startToken [r]
acc ReadState
state ByteString
bytes = do
        cols <- IORef (Maybe ResultSetInfo) -> IO (Maybe ResultSetInfo)
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef (Maybe ResultSetInfo)
arResultSet ActiveResponse
ar)
        let token = case RowRead r
reader of DrainRows r
_ -> Maybe ResultSetInfo -> Parser Token
readTokenSkipping Maybe ResultSetInfo
cols; RowRead r
_ -> Int -> Maybe ResultSetInfo -> Parser Token
readTokenWith Int
ahead Maybe ResultSetInfo
cols
            -- Rows and output parameters cost their values' size.  Every other
            -- token is metadata, whose lengths the bytes alone don't bound.
            recipe = if ByteString -> Bool
isRowToken ByteString
bytes Bool -> Bool -> Bool
|| ByteString -> Bool
isReturnValueToken ByteString
bytes then Parser Token
token else Int -> Parser Token -> Parser Token
forall a. Int -> Parser a -> Parser a
P.limited Int
tokenSizeLimit Parser Token
token
        resume acc state (fmap (\(P.Got Token
tok ByteString
rest) -> Token -> ByteString -> ReadEvent
ReadToken Token
tok ByteString
rest) (P.parse recipe bytes))
      accept :: [r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc Bool
final ByteString
rest = \case
        Left ReadFailure
failure -> [r]
-> Bool -> ByteString -> Int -> ReadFailure -> IO (StepProgress r)
failedRow [r]
acc Bool
final ByteString
rest Int
0 ReadFailure
failure
        Right r
r -> [r] -> Bool -> ByteString -> Int -> IO (StepProgress r)
atBytes (r
r r -> [r] -> [r]
forall a. a -> [a] -> [a]
: [r]
acc) Bool
final ByteString
rest Int
0
      driveRow :: [r]
-> ReadState
-> IORef (Maybe (Either ReadFailure r))
-> Stream Await IO (Got (Either ReadFailure r))
-> IO (StepProgress r)
driveRow [r]
acc ReadState
state IORef (Maybe (Either ReadFailure r))
inbox Stream Await IO (Got (Either ReadFailure r))
stream = Stream Await IO (Got (Either ReadFailure r))
-> IO
     (Either
        (Got (Either ReadFailure r))
        (Await (Stream Await IO (Got (Either ReadFailure r)))))
forall (m :: * -> *) (f :: * -> *) r.
Monad m =>
Stream f m r -> m (Either r (f (Stream f m r)))
inspect Stream Await IO (Got (Either ReadFailure r))
stream IO
  (Either
     (Got (Either ReadFailure r))
     (Await (Stream Await IO (Got (Either ReadFailure r)))))
-> (Either
      (Got (Either ReadFailure r))
      (Await (Stream Await IO (Got (Either ReadFailure r))))
    -> IO (StepProgress r))
-> IO (StepProgress r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Left (P.Got Either ReadFailure r
result ByteString
rest) -> [r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
rest Either ReadFailure r
result
        Right Await (Stream Await IO (Got (Either ReadFailure r)))
await -> [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [r]
acc ReadState
state (Stream Await IO ReadEvent -> IO (StepProgress r))
-> Stream Await IO ReadEvent -> IO (StepProgress r)
forall a b. (a -> b) -> a -> b
$
          (Got (Either ReadFailure r) -> ReadEvent)
-> Stream Await IO (Got (Either ReadFailure r))
-> Stream Await IO ReadEvent
forall a b. (a -> b) -> Stream Await IO a -> Stream Await IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(P.Got Either ReadFailure r
result ByteString
rest) -> IO () -> ByteString -> ReadEvent
ReadRow (IORef (Maybe (Either ReadFailure r))
-> Maybe (Either ReadFailure r) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (Either ReadFailure r))
inbox (Either ReadFailure r -> Maybe (Either ReadFailure r)
forall a. a -> Maybe a
Just Either ReadFailure r
result)) ByteString
rest) (Await (Stream Await IO (Got (Either ReadFailure r)))
-> Stream Await IO (Got (Either ReadFailure r))
forall (m :: * -> *) (f :: * -> *) r.
(Monad m, Functor f) =>
f (Stream f m r) -> Stream f m r
Streaming.wrap Await (Stream Await IO (Got (Either ReadFailure r)))
await)
      resume :: [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [r]
acc ReadState
state Stream Await IO ReadEvent
suspended = Stream Await IO ReadEvent
-> IO (Either ReadEvent (Await (Stream Await IO ReadEvent)))
forall (m :: * -> *) (f :: * -> *) r.
Monad m =>
Stream f m r -> m (Either r (f (Stream f m r)))
inspect Stream Await IO ReadEvent
suspended IO (Either ReadEvent (Await (Stream Await IO ReadEvent)))
-> (Either ReadEvent (Await (Stream Await IO ReadEvent))
    -> IO (StepProgress r))
-> IO (StepProgress r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Left ReadEvent
event -> case ReadEvent
event of
          ReadToken Token
tok ByteString
rest -> do
            Token -> IO ()
applyToken Token
tok
            let after :: ReadState
after = ReadState
state { readPosition = AtToken rest }
            case Token
tok of
              TRow Row
row -> case RowRead r
reader of
                WholeRows Row -> r
decode -> [r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
rest (r -> Either ReadFailure r
forall a b. b -> Either a b
Right (Row -> r
decode Row
row))
                RowRead r
_ -> IO (StepProgress r)
-> (r -> IO (StepProgress r)) -> Maybe r -> IO (StepProgress r)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO (StepProgress r)
forall {b}. IO b
changed ([r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
rest (Either ReadFailure r -> IO (StepProgress r))
-> (r -> Either ReadFailure r) -> r -> IO (StepProgress r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Either ReadFailure r
forall a b. b -> Either a b
Right) (RowRead r -> Maybe r
forall a. RowRead a -> Maybe a
skippedAs RowRead r
reader)
              Token
_ -> IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
after IO () -> IO (StepProgress r) -> IO (StepProgress r)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> StepProgress r -> IO (StepProgress r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResponseStep r -> StepProgress r
forall r. ResponseStep r -> StepProgress r
StepReady (Token -> ResponseStep r
forall r. Token -> ResponseStep r
StepToken Token
tok))
          ReadRow IO ()
deliver ByteString
rest -> do
            case RowRead r -> Maybe r
forall a. RowRead a -> Maybe a
skippedAs RowRead r
reader of
              Just r
r -> [r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
rest (r -> Either ReadFailure r
forall a b. b -> Either a b
Right r
r)
              Maybe r
Nothing -> case RowRead r
reader of
                ReadRows Parser r
_ IORef (Maybe (Either ReadFailure r))
inbox -> do
                  IO ()
deliver
                  result <- IORef (Maybe (Either ReadFailure r))
-> IO (Maybe (Either ReadFailure r))
forall a. IORef a -> IO a
readIORef IORef (Maybe (Either ReadFailure r))
inbox
                  writeIORef inbox Nothing
                  maybe changed (accept acc (readFinalPacket state) rest) result
                RowRead r
_ -> IO (StepProgress r)
forall {b}. IO b
changed
          SkippedRow ByteString
rest -> do
            let after :: ReadState
after = ReadState
state { readPosition = AtToken rest }
            IO (StepProgress r)
-> (r -> IO (StepProgress r)) -> Maybe r -> IO (StepProgress r)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([r] -> ReadState -> IO (StepProgress r)
at [r]
acc ReadState
after) ([r]
-> Bool
-> ByteString
-> Either ReadFailure r
-> IO (StepProgress r)
accept [r]
acc (ReadState -> Bool
readFinalPacket ReadState
state) ByteString
rest (Either ReadFailure r -> IO (StepProgress r))
-> (r -> Either ReadFailure r) -> r -> IO (StepProgress r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Either ReadFailure r
forall a b. b -> Either a b
Right) (RowRead r -> Maybe r
forall a. RowRead a -> Maybe a
skippedAs RowRead r
reader)
        Right await :: Await (Stream Await IO ReadEvent)
await@(P.NeedMore Feed -> Stream Await IO ReadEvent
k) -> do
          -- Committed before the receive, the only place this step can be
          -- interrupted.  A saved step is run once, which a value's accumulator
          -- relies on: it writes into a buffer the suspension holds.
          let pending :: ReadState
pending = ReadState
state { readPosition = Reading (Streaming.wrap await) }
          IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
pending
          if Bool -> Bool
not ([r] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [r]
acc) then ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
pending [r]
acc Maybe DecodeError
forall a. Maybe a
Nothing else
            if ReadState -> Bool
readFinalPacket ReadState
state then [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [] ReadState
pending (Feed -> Stream Await IO ReadEvent
k Feed
P.EndOfMessage) else do
              (body, final) <- Connection -> Session -> IO (ByteString, Bool)
readResponsePacket Connection
conn Session
sess
              -- 'P.Discard' lets a value being assembled drop its buffer, since
              -- nothing will read it.
              resume [] pending { readFinalPacket = final } (k (case reader of DrainRows r
_ -> ByteString -> Feed
P.Discard ByteString
body; RowRead r
_ -> ByteString -> Feed
P.More ByteString
body))
        Right await :: Await (Stream Await IO ReadEvent)
await@(P.RunConsumer Bool
live Runner -> IO x
io Maybe x -> Stream Await IO ReadEvent
k) -> do
          -- Committed before the consumer runs.  If it throws, a drain resumes
          -- here, and a drain gives 'Nothing', which reads past the rest of the
          -- value without running the consumer again.
          let pending :: ReadState
pending = ReadState
state { readPosition = Reading (Streaming.wrap await) }
          IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
pending
          -- A consumer that pulls its value from the response may wait for the
          -- server, so the complete rows go back first, as before a receive; the
          -- next step runs it.  One given its value whole runs here, among them.
          if Bool
live Bool -> Bool -> Bool
&& Bool -> Bool
not ([r] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [r]
acc) then ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
forall {r}.
ReadState -> [r] -> Maybe DecodeError -> IO (StepProgress r)
rows ReadState
pending [r]
acc Maybe DecodeError
forall a. Maybe a
Nothing else case RowRead r -> Maybe r
forall a. RowRead a -> Maybe a
skippedAs RowRead r
reader of
            -- A drain must not run user code.
            Just r
_ -> [r]
-> ReadState -> Stream Await IO ReadEvent -> IO (StepProgress r)
resume [r]
acc ReadState
state (Maybe x -> Stream Await IO ReadEvent
k Maybe x
forall a. Maybe a
Nothing)
            Maybe r
Nothing -> StepProgress r -> IO (StepProgress r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StepProgress r -> IO (StepProgress r))
-> StepProgress r -> IO (StepProgress r)
forall a b. (a -> b) -> a -> b
$ IO x -> (x -> IO (StepProgress r)) -> StepProgress r
forall r x. IO x -> (x -> IO (StepProgress r)) -> StepProgress r
StepConsumer ((Runner -> IO x) -> IO x
forall a. (Runner -> IO a) -> IO a
runConsumer Runner -> IO x
io) ((x -> IO (StepProgress r)) -> StepProgress r)
-> (x -> IO (StepProgress r)) -> StepProgress r
forall a b. (a -> b) -> a -> b
$ \x
result -> do
              final <- ReadState -> Bool
readFinalPacket (ReadState -> Bool) -> IO ReadState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef ReadState -> IO ReadState
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar)
              resume acc state { readFinalPacket = final } (k (Just result))
      -- An exception from the consumer ends the query at once, as one from the
      -- caller's own code would, and leaves the rest of the value and row to the
      -- drain.  Reading them here would make a timeout in a consumer wait for the
      -- server, and hang if the server had stalled.  The statement is canceled
      -- and the handle abandoned, so that a stream of the response can't run the
      -- consumer again.
      runConsumer :: (P.Runner -> IO a) -> IO a
      runConsumer :: forall a. (Runner -> IO a) -> IO a
runConsumer Runner -> IO a
io = IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ IO ()
enter IO ()
leave (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
        -- The consumer's pulls run back under the mask, one packet at a time.  A
        -- consumer could catch a failure of the byte stream and return, so each
        -- pull breaks the connection itself.
        let feed :: IO Feed
feed = do
              final <- ReadState -> Bool
readFinalPacket (ReadState -> Bool) -> IO ReadState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef ReadState -> IO ReadState
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar)
              if final then pure P.EndOfMessage else do
                (body, end) <- readResponsePacket conn sess
                modifyIORef' (arState ar) (\ReadState
s -> ReadState
s { readFinalPacket = end })
                pure (P.More body)
            pull :: Runner
pull = IO Feed -> Runner
P.feedRunner IO Feed
feed
            runner :: Runner
runner = (forall a. IORef (Stream Await IO (Got a)) -> IO (Got a)) -> Runner
P.Runner ((forall a. IORef (Stream Await IO (Got a)) -> IO (Got a))
 -> Runner)
-> (forall a. IORef (Stream Await IO (Got a)) -> IO (Got a))
-> Runner
forall a b. (a -> b) -> a -> b
$ \IORef (Stream Await IO (Got a))
current -> IO (Got a) -> IO (Got a)
forall a. IO a -> IO a
mask_ (IO (Got a) -> IO (Got a)) -> IO (Got a) -> IO (Got a)
forall a b. (a -> b) -> a -> b
$
              Connection -> IO (Got a) -> IO (Got a)
forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn (Runner -> forall a. IORef (Stream Await IO (Got a)) -> IO (Got a)
P.runAwait Runner
pull IORef (Stream Await IO (Got a))
current)
        IO a -> IO a
forall a. IO a -> IO a
restore (Runner -> IO a
io Runner
runner) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` Connection -> ActiveResponse -> IO ()
abandonAtOnce Connection
conn ActiveResponse
ar
      enter :: IO ()
enter = do
        IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef Bool
arInConsumer ActiveResponse
ar) Bool
True
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe Mars -> Bool
forall a. Maybe a -> Bool
isNothing (Connection -> Maybe Mars
connMars Connection
conn)) (IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
connInCallback Connection
conn) Bool
True)
      leave :: IO ()
leave = do
        IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef Bool
arInConsumer ActiveResponse
ar) Bool
False
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe Mars -> Bool
forall a. Maybe a -> Bool
isNothing (Connection -> Maybe Mars
connMars Connection
conn)) (IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
connInCallback Connection
conn) Bool
False)
      messageEnd :: IO (StepProgress r)
messageEnd = do
        -- After an Attention the acknowledgement may come in a message of its
        -- own, so the end of this one doesn't end the response.
        cancellation <- IORef Cancellation -> IO Cancellation
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef Cancellation
arCancellation ActiveResponse
ar)
        final <- readIORef (arSawFinalDone ar)
        if cancellation == AttentionSent
          then writeIORef (arState ar) (ReadState False (AtToken B.empty)) >> loop []
          else if not final && cancellation /= AttentionAcknowledged
            then throwIO (ProtocolError "response message ended without a final DONE token")
            else writeIORef (sessResponse sess) Nothing >> pure (StepReady StepEnd)
      applyToken :: Token -> IO ()
applyToken Token
tok = do
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Token -> Bool
isFinalDone Token
tok) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef Bool
arSawFinalDone ActiveResponse
ar) Bool
True
        Maybe Done -> (Done -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Token -> Maybe Done
doneOf Token
tok) ((Done -> IO ()) -> IO ()) -> (Done -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Done
d -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Done -> Bool
doneAttention Done
d) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          IORef Cancellation -> (Cancellation -> Cancellation) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (ActiveResponse -> IORef Cancellation
arCancellation ActiveResponse
ar) (\Cancellation
s -> if Cancellation
s Cancellation -> Cancellation -> Bool
forall a. Eq a => a -> a -> Bool
== Cancellation
AttentionSent then Cancellation
AttentionAcknowledged else Cancellation
s)
        case Token
tok of
          TError ServerMessage
msg | ServerMessage -> Word8
msgClass ServerMessage
msg Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
20 -> do
            let err :: SqlServerError
err = ServerMessage -> Maybe Text -> SqlServerError
SqlServerError ServerMessage
msg (ActiveResponse -> Maybe Text
arStatement ActiveResponse
ar)
            Connection -> UnusableReason -> IO ()
markBroken Connection
conn (ServerMessage -> UnusableReason
FatalServerError ServerMessage
msg)
            SqlServerError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SqlServerError
err
          TColMetaData [ColumnInfo]
cols -> IORef (Maybe ResultSetInfo) -> Maybe ResultSetInfo -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef (Maybe ResultSetInfo)
arResultSet ActiveResponse
ar) (ResultSetInfo -> Maybe ResultSetInfo
forall a. a -> Maybe a
Just ([ColumnInfo] -> ResultSetInfo
compileResultSet [ColumnInfo]
cols))
          TEnvChange EnvChange
ec -> do
            IORef TdsContext -> (TdsContext -> TdsContext) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef TdsContext
connContext Connection
conn) (EnvChange -> TdsContext -> TdsContext
applyEnvChange EnvChange
ec)
            case EnvChange
ec of EnvChange
ECResetConnectionAck -> IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef Bool
arSawResetAck ActiveResponse
ar) Bool
True; EnvChange
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Token
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  -- One handler for each stretch of reading, not one for each completed row.  It
  -- also preserves failures of custom pure conversions.
  let drive :: IO (StepProgress r) -> IO (ResponseStep r)
drive IO (StepProgress r)
reading = Connection -> IO (StepProgress r) -> IO (StepProgress r)
forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn IO (StepProgress r)
reading IO (StepProgress r)
-> (StepProgress r -> IO (ResponseStep r)) -> IO (ResponseStep r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        StepReady ResponseStep r
step -> ResponseStep r -> IO (ResponseStep r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResponseStep r
step
        StepConsumer IO x
consumer x -> IO (StepProgress r)
continue -> IO x
consumer IO x -> (x -> IO (ResponseStep r)) -> IO (ResponseStep r)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (StepProgress r) -> IO (ResponseStep r)
drive (IO (StepProgress r) -> IO (ResponseStep r))
-> (x -> IO (StepProgress r)) -> x -> IO (ResponseStep r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> IO (StepProgress r)
continue
  IO (StepProgress r) -> IO (ResponseStep r)
forall {r}. IO (StepProgress r) -> IO (ResponseStep r)
drive ([r] -> IO (StepProgress r)
loop [])

-- | The next payload of this session's response, and whether it is the
-- message's last.  This hides the transport, TLS and SMP from the parser, and
-- checks the packet's type.
readResponsePacket :: Connection -> Session -> IO (B.ByteString, Bool)
readResponsePacket :: Connection -> Session -> IO (ByteString, Bool)
readResponsePacket Connection
conn Session
sess = Connection -> IO (ByteString, Bool) -> IO (ByteString, Bool)
forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn (IO (ByteString, Bool) -> IO (ByteString, Bool))
-> IO (ByteString, Bool) -> IO (ByteString, Bool)
forall a b. (a -> b) -> a -> b
$ do
  (header, bytes) <- Connection -> Session -> IO (Header, ByteString)
readSessionPacket Connection
conn Session
sess
  unless (headerType header == PTTabularResult) $
    throwIO (ProtocolError ("unexpected packet type " ++ show (headerType header)))
  pure (bytes, headerStatus header `hasStatus` statusEndOfMessage)

-- | A consumer or cell pull cannot resume after an exception, unlike a plain
-- read after an interruption, so its handle is abandoned even if the response
-- has fully arrived. The caller holds the connection lock; the saved recipe
-- remains available to its drain.
abandonAtOnce :: Connection -> ActiveResponse -> IO ()
abandonAtOnce :: Connection -> ActiveResponse -> IO ()
abandonAtOnce Connection
conn ActiveResponse
ar = do
  -- A consumer's pull that broke the connection fails through the consumer.
  broken <- Maybe UnusableReason -> Bool
forall a. Maybe a -> Bool
isJust (Maybe UnusableReason -> Bool)
-> IO (Maybe UnusableReason) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Maybe UnusableReason) -> IO (Maybe UnusableReason)
forall a. IORef a -> IO a
readIORef (Connection -> IORef (Maybe UnusableReason)
connBroken Connection
conn)
  unless broken $ void (cancelStatement conn ar)
  markAbandoned ar
  abandonHandle ar

-- | Pull a generic cell fragment, preserving a recipe for abandoning its row.
-- The function says, from what was read, how to skip the rest of the row, or
-- 'Nothing' when the pull ended at a token boundary.
readPiece :: Connection -> ActiveResponse -> P.Parser a -> (a -> Maybe (P.Parser ())) -> IO a
readPiece :: forall a.
Connection
-> ActiveResponse -> Parser a -> (a -> Maybe (Parser ())) -> IO a
readPiece Connection
conn ActiveResponse
ar Parser a
recipe a -> Maybe (Parser ())
remainder = Connection -> IO a -> IO a
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ IO a -> IO a
forall a. IO a -> IO a
mask_ (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
  Connection -> ActiveResponse -> IO ()
checkCurrent Connection
conn ActiveResponse
ar
  state <- IORef ReadState -> IO ReadState
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar)
  let start ByteString
bytes = ReadState -> Stream Await IO (Got a) -> IO a
drive ReadState
state (Parser a -> ByteString -> Stream Await IO (Got a)
forall a. Parser a -> ByteString -> Stream Await IO (Got a)
P.parse Parser a
recipe ByteString
bytes)
      drive ReadState
current Stream Await IO (Got a)
stream = Connection
-> IO (Either (Got a) (Await (Stream Await IO (Got a))))
-> IO (Either (Got a) (Await (Stream Await IO (Got a))))
forall a. Connection -> IO a -> IO a
breakOnFailure Connection
conn (Stream Await IO (Got a)
-> IO (Either (Got a) (Await (Stream Await IO (Got a))))
forall (m :: * -> *) (f :: * -> *) r.
Monad m =>
Stream f m r -> m (Either r (f (Stream f m r)))
inspect Stream Await IO (Got a)
stream) IO (Either (Got a) (Await (Stream Await IO (Got a))))
-> (Either (Got a) (Await (Stream Await IO (Got a))) -> IO a)
-> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Left (P.Got a
a ByteString
rest) -> do
          IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
current { readPosition = maybe (AtToken rest) (Parked rest) (remainder a) }
          a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
        Right await :: Await (Stream Await IO (Got a))
await@(P.NeedMore Feed -> Stream Await IO (Got a)
k) -> do
          -- If this wait is interrupted, no later pull will come for this piece,
          -- so what is saved is the piece followed by the skip of its row.
          let drain :: Stream Await IO ReadEvent
drain = Await (Stream Await IO (Got a)) -> Stream Await IO (Got a)
forall (m :: * -> *) (f :: * -> *) r.
(Monad m, Functor f) =>
f (Stream f m r) -> Stream f m r
Streaming.wrap Await (Stream Await IO (Got a))
await Stream Await IO (Got a)
-> (Got a -> Stream Await IO ReadEvent)
-> Stream Await IO ReadEvent
forall a b.
Stream Await IO a -> (a -> Stream Await IO b) -> Stream Await IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(P.Got a
a ByteString
rest) -> case a -> Maybe (Parser ())
remainder a
a of
                Maybe (Parser ())
Nothing -> ReadEvent -> Stream Await IO ReadEvent
forall a. a -> Stream Await IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> ReadEvent
SkippedRow ByteString
rest)
                Just Parser ()
skip -> (Got () -> ReadEvent)
-> Stream Await IO (Got ()) -> Stream Await IO ReadEvent
forall a b. (a -> b) -> Stream Await IO a -> Stream Await IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(P.Got () ByteString
rest') -> ByteString -> ReadEvent
SkippedRow ByteString
rest') (Parser () -> ByteString -> Stream Await IO (Got ())
forall a. Parser a -> ByteString -> Stream Await IO (Got a)
P.parse Parser ()
skip ByteString
rest)
              pending :: ReadState
pending = ReadState
current { readPosition = Reading drain }
          IORef ReadState -> ReadState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (ActiveResponse -> IORef ReadState
arState ActiveResponse
ar) ReadState
pending
          if ReadState -> Bool
readFinalPacket ReadState
current then ReadState -> Stream Await IO (Got a) -> IO a
drive ReadState
pending (Feed -> Stream Await IO (Got a)
k Feed
P.EndOfMessage) else do
            (bytes, final) <- Connection -> Session -> IO (ByteString, Bool)
readResponsePacket Connection
conn (ActiveResponse -> Session
arSession ActiveResponse
ar)
            drive pending { readFinalPacket = final } (k (P.More bytes))
        Right P.RunConsumer{} -> ProtocolError -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ([Char] -> ProtocolError
ProtocolError [Char]
"a cell pull cannot invoke a field consumer")
      interrupted (SomeAsyncException
err :: SomeAsyncException) = Connection -> ActiveResponse -> IO ()
abandonAtOnce Connection
conn ActiveResponse
ar IO () -> IO b -> IO b
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeAsyncException -> IO b
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeAsyncException
err
  (case readPosition state of
    AtToken ByteString
bytes -> ByteString -> IO a
start ByteString
bytes
    Parked ByteString
bytes Parser ()
_ -> ByteString -> IO a
start ByteString
bytes
    -- A pull that was interrupted left only the drain above.
    Reading Stream Await IO ReadEvent
_ -> ResultAbandoned -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO ResultAbandoned
ResultAbandoned) `catch` interrupted

-- | Whether the next token starts a row.
isRowToken :: B.ByteString -> Bool
isRowToken :: ByteString -> Bool
isRowToken ByteString
bs = ByteString -> Int -> Bool
isRowTokenAt ByteString
bs Int
0

-- | Whether the token at this position of the payload starts a row.
isRowTokenAt :: B.ByteString -> Int -> Bool
isRowTokenAt :: ByteString -> Int -> Bool
isRowTokenAt ByteString
payload Int
pos = Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ByteString -> Int
B.length ByteString
payload Bool -> Bool -> Bool
&& (Word8
byte Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0xD1 Bool -> Bool -> Bool
|| Word8
byte Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0xD2)
  where byte :: Word8
byte = HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
B.index ByteString
payload Int
pos

isReturnValueToken :: B.ByteString -> Bool
isReturnValueToken :: ByteString -> Bool
isReturnValueToken ByteString
bytes = Bool -> Bool
not (ByteString -> Bool
B.null ByteString
bytes) Bool -> Bool -> Bool
&& HasCallStack => ByteString -> Word8
ByteString -> Word8
B.head ByteString
bytes Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0xAC

applyEnvChange :: EnvChange -> TdsContext -> TdsContext
applyEnvChange :: EnvChange -> TdsContext -> TdsContext
applyEnvChange EnvChange
ec TdsContext
ctx = case EnvChange
ec of
  ECPacketSize{Word32
newSize :: Word32
newSize :: EnvChange -> Word32
newSize} -> TdsContext
ctx { tdsPacketSize = newSize }
  ECBeginTransaction{TransactionDescriptor
transaction :: TransactionDescriptor
transaction :: EnvChange -> TransactionDescriptor
transaction} -> TdsContext
ctx { tdsTransaction = Just transaction }
  ECCommitTransaction{} -> TdsContext
ctx { tdsTransaction = Nothing }
  ECRollbackTransaction{} -> TdsContext
ctx { tdsTransaction = Nothing }
  ECDefectTransaction{} -> TdsContext
ctx { tdsTransaction = Nothing }
  ECTransactionEnded{} -> TdsContext
ctx { tdsTransaction = Nothing }
  ECEnlistDtcTransaction{TransactionDescriptor
transaction :: EnvChange -> TransactionDescriptor
transaction :: TransactionDescriptor
transaction} -> TdsContext
ctx { tdsTransaction = Just transaction }
  -- A reset rolls back any open transaction.
  EnvChange
ECResetConnectionAck -> TdsContext
ctx { tdsTransaction = Nothing }
  -- A copy: the bytes are a slice of whatever buffer the response was
  -- received into, and the context lives as long as the connection.  As
  -- a slice, the collation of a plaintext login kept the 32 KiB buffer of
  -- that login's read alive ('recvFresh').
  -- (Forced, or the copy would be a thunk that holds the slice.)
  ECSqlCollation{ByteString
newBytes :: ByteString
newBytes :: EnvChange -> ByteString
newBytes} -> TdsContext
ctx { tdsCollation = if B.null newBytes then Nothing else Just $! B.copy newBytes }
  EnvChange
_ -> TdsContext
ctx

-- | Read and discard the rest of the current response, a batch of rows
-- at a time (nothing is decoded).  The record given is the session's
-- current one (read from 'sessResponse', as 'finishSession' does), never
-- a caller's handle: a handle abandoned by an interruption has an older
-- generation, which 'stepResponse' would refuse.  ("Hmssql.BatchQuery"
-- drains with its reader's record on the synchronous path only, where no
-- generation has changed.)
drainActiveResponse :: Connection -> ActiveResponse -> IO ()
drainActiveResponse :: Connection -> ActiveResponse -> IO ()
drainActiveResponse Connection
conn ActiveResponse
ar = do
  r <- Connection -> ActiveResponse -> RowRead () -> IO (ResponseStep ())
forall r.
Connection -> ActiveResponse -> RowRead r -> IO (ResponseStep r)
nextStep Connection
conn ActiveResponse
ar RowRead ()
drainingRows
  case r of
    ResponseStep ()
StepEnd -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    ResponseStep ()
_ -> Connection -> ActiveResponse -> IO ()
drainActiveResponse Connection
conn ActiveResponse
ar

-- | Make sure no response is pending on the connection: 'finishSession'
-- for its session, or under MARS for every session that has one.  Under
-- MARS the Attentions all go out first and the acknowledgements are read
-- after, so settling several running statements costs one round trip, not
-- one each.
finishResponse :: Connection -> IO ()
finishResponse :: Connection -> IO ()
finishResponse Connection
conn = Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  sessions <- Connection -> IO [Session]
allSessions Connection
conn
  when (isJust (connMars conn)) $
    forM_ sessions $ \Session
sess -> IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ActiveResponse -> IO ()) -> Maybe ActiveResponse -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Connection -> ActiveResponse -> IO ()
cancelRunning Connection
conn)
  mapM_ (finishSession conn) sessions

-- | Make sure no response is pending on the session: send an Attention if
-- the server is still sending (unless one was sent already, by
-- 'nextToken'), then discard everything up to the end of the message (and,
-- after an Attention, up to its acknowledgement).  A response whose last
-- packet has arrived ('fullyArrived') is read to its end with nothing
-- sent.  Returns once the response is over on the server too.
--
-- If the drain fails or is interrupted, the connection is marked broken:
-- how much of the canceled response has been read is then unknown, and a
-- second Attention would get a second acknowledgement, which would be read
-- as the reply to the next request.
finishSession :: Connection -> Session -> IO ()
finishSession :: Connection -> Session -> IO ()
finishSession Connection
conn Session
sess = Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
  IORef (Maybe ActiveResponse) -> IO (Maybe ActiveResponse)
forall a. IORef a -> IO a
readIORef (Session -> IORef (Maybe ActiveResponse)
sessResponse Session
sess) IO (Maybe ActiveResponse)
-> (Maybe ActiveResponse -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe ActiveResponse
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just ActiveResponse
ar -> do
      Connection -> ActiveResponse -> IO ()
cancelRunning Connection
conn ActiveResponse
ar
      Connection -> ActiveResponse -> IO ()
drainActiveResponse Connection
conn ActiveResponse
ar IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
e :: SomeException) -> do
        Connection -> UnusableReason -> IO ()
markBroken Connection
conn ((SomeException -> UnusableReason)
-> SomeException -> UnusableReason
reasonFor SomeException -> UnusableReason
DrainFailed SomeException
e)
        SomeException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e
      -- A canceled request that carried a reset: SQL Server applies the
      -- reset all the same, and acknowledges it in the canceled response
      -- (every time, in 6000 tries: batch, RPC and TM requests, the
      -- Attention sent 0 to 3 ms behind them, with and without MARS).  If
      -- a response ever ends without the acknowledgement, the reset is
      -- asked for again.
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ActiveResponse -> Bool
arCarriedReset ActiveResponse
ar) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        acked <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (ActiveResponse -> IORef Bool
arSawResetAck ActiveResponse
ar)
        unless acked $ atomicWriteIORef (connResetPending conn) True

-- | Send an Attention on the response's session, unless the response has
-- fully arrived ('fullyArrived': the statement is over, and an Attention
-- would get an acknowledgement of its own) or one was sent already.
-- Masked so the Attention is recorded once it was sent.  The caller
-- holds the lock and drains the response next ('finishSession').
cancelRunning :: Connection -> ActiveResponse -> IO ()
cancelRunning :: Connection -> ActiveResponse -> IO ()
cancelRunning Connection
conn ActiveResponse
ar = do
  arrived <- ActiveResponse -> IO Bool
fullyArrived ActiveResponse
ar
  attentionSent <- attentionWasSent <$> readIORef (arCancellation ar)
  unless (arrived || attentionSent) $ mask_ $ do
    ctx <- readIORef (connContext conn)
    sendFramed conn (arSession ar) ctx PTAttention (PacketStatus 0) LB.empty
    writeIORef (arCancellation ar) AttentionSent

-- | Cancel one response and leave the others alone: 'finishSession' for
-- its session, if the response is still that session's (read part-way, or
-- canceled by an interrupted read and not yet drained); nothing if it has
-- ended, was canceled and drained already, or the connection is closed or
-- broken.  A stream of the response throws 'ResultAbandoned' if read on.
cancelActiveResponse :: Connection -> ActiveResponse -> IO ()
cancelActiveResponse :: Connection -> ActiveResponse -> IO ()
cancelActiveResponse Connection
conn ActiveResponse
ar = Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
withConnectionLock Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  ended <- Maybe SomeException -> Bool
forall a. Maybe a -> Bool
isJust (Maybe SomeException -> Bool)
-> IO (Maybe SomeException) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Maybe SomeException) -> IO (Maybe SomeException)
forall a. TVar a -> IO a
readTVarIO (Connection -> TVar (Maybe SomeException)
connEnded Connection
conn)
  broken <- isJust <$> readIORef (connBroken conn)
  current <- readIORef (sessResponse sess)
  case current of
    -- The same response, whatever its generation (an interrupted read
    -- gives it a new one).
    Just ActiveResponse
ar' | Bool -> Bool
not (Bool
ended Bool -> Bool -> Bool
|| Bool
broken), ActiveResponse -> IORef ReadState
arState ActiveResponse
ar' IORef ReadState -> IORef ReadState -> Bool
forall a. Eq a => a -> a -> Bool
== ActiveResponse -> IORef ReadState
arState ActiveResponse
ar -> Connection -> Session -> IO ()
finishSession Connection
conn Session
sess
    Maybe ActiveResponse
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    sess :: Session
sess = ActiveResponse -> Session
arSession ActiveResponse
ar

-- | The SQL text of each response on the connection that has not been
-- read to its end or canceled, where known.  The caller holds the lock.
busyStatements :: Connection -> IO [Text]
busyStatements :: Connection -> IO [Text]
busyStatements Connection
conn = do
  sessions <- Connection -> IO [Session]
allSessions Connection
conn
  busy <- catMaybes <$> forM sessions (readIORef . sessResponse)
  pure [fromMaybe "(a request without SQL text)" (arStatement ar) | ar <- busy]