Transport.hs

{-# LANGUAGE CPP #-}

-- | The byte-level transport: a TCP socket with a read buffer, plus the framing
-- that turns the stream into discrete @[type byte][Int32 length][body]@
-- messages.
module Pqi.Native.Transport
  ( Transport,
    connect,
    isUnixSocketHost,
    unixSocketPath,
    unconnected,
    close,
    send,
    receiveFrame,
    socketFd,
    peerIp,
    readUntilClosed,
  )
where

import Control.Exception (mask_)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as ByteString.Char8
import qualified Data.ByteString.Internal as ByteString.Internal
import Data.IORef
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (plusPtr)
#if !defined(mingw32_HOST_OS)
import GHC.Conc (atomically, orElse, threadWaitReadSTM, threadWaitWriteSTM)
#endif
import qualified Network.Socket as Socket
import qualified Network.Socket.ByteString as Socket.ByteString
import Pqi.Native.Transport.Prelude
import qualified PtrPoker.Write as Poker
import System.IO.Error (eofErrorType, mkIOError)

-- | An open connection's byte transport.
data Transport = Transport
  { socket :: Socket.Socket,
    readBuffer :: IORef ReadBuffer
  }

-- | Bytes received from the socket and not yet consumed: the region
-- @[start, end)@ of a pinned @storage@ of @capacity@ bytes. The socket is
-- read straight into the free region @[end, capacity)@, so a received byte
-- is copied only when the live region is moved to a larger storage.
--
-- Bytes before @end@ are never written again: the frames 'receiveFrame'
-- hands out are slices of the storage, so making room always allocates a
-- fresh storage rather than compacting in place.
data ReadBuffer = ReadBuffer
  { storage :: ForeignPtr Word8,
    capacity :: Int,
    start :: Int,
    end :: Int
  }

newReadBuffer :: IO (IORef ReadBuffer)
newReadBuffer = do
  storage <- ByteString.Internal.mallocByteString 0
  newIORef ReadBuffer {storage, capacity = 0, start = 0, end = 0}

-- | The buffered bytes, as a slice of the storage.
buffered :: ReadBuffer -> ByteString
buffered buffer =
  ByteString.Internal.fromForeignPtr (storage buffer) (start buffer) (end buffer - start buffer)

-- | Make room for the buffer to hold @n@ bytes in all, moving the buffered
-- bytes to a fresh storage if the current one cannot. The new storage is
-- exactly @n@ bytes when that is what a large frame needs, and otherwise
-- at least 'minimumCapacity' and twice the bytes being moved, so that a
-- buffer grown one receive at a time copies each byte a bounded number of
-- times.
reserve :: Int -> ReadBuffer -> IO ReadBuffer
reserve n buffer
  | start buffer + n <= capacity buffer = pure buffer
  | otherwise = do
      let live = end buffer - start buffer
          newCapacity = maximum [n, minimumCapacity, 2 * live]
      newStorage <- ByteString.Internal.mallocByteString newCapacity
      withForeignPtr (storage buffer) \source ->
        withForeignPtr newStorage \target ->
          copyBytes target (source `plusPtr` start buffer) live
      pure ReadBuffer {storage = newStorage, capacity = newCapacity, start = 0, end = live}

-- | The smallest storage allocated, which is also the most a single receive
-- reads while only small messages are being waited for.
minimumCapacity :: Int
minimumCapacity = 65536

-- | Receive once from the socket into the buffer's free region, after making
-- room for @n@ bytes in all, reporting whether the peer closed the
-- connection. Blocks until bytes are available.
--
-- The receive and the recording of its bytes happen under 'mask_', so an
-- async exception can never land between @recv@ returning and its bytes
-- being recorded, which would drop them and desync the stream. @recv@
-- itself stays interruptible inside 'mask_'. Making room is a pure
-- replacement of the buffer, safe to interrupt at any point.
receiveSome :: Transport -> Int -> IO Bool
receiveSome transport n = do
  buffer <- readIORef (readBuffer transport) >>= reserve n
  writeIORef (readBuffer transport) buffer
  mask_ do
    received <-
      withForeignPtr (storage buffer) \pointer ->
        Socket.recvBuf (socket transport) (pointer `plusPtr` end buffer) (capacity buffer - end buffer)
    if received == 0
      then pure True
      else do
        writeIORef (readBuffer transport) buffer {end = end buffer + received}
        pure False

-- | Open a connection to the given host, numeric host address and port, in
-- the order libpq considers them:
--
-- 1. A non-empty @hostaddr@ is connected to over TCP, and the @host@ is not
--    consulted at all: it is not resolved, and a path-like one does not
--    select a Unix-domain socket. libpq keeps it only for authentication and
--    for naming the server in messages.
-- 2. Otherwise, if the host looks like an absolute path (starts with @\/@ -
--    see 'isUnixSocketHost'), it names a Unix-domain socket /directory/
--    rather than a TCP host, and the connection is made to the socket file
--    @'unixSocketPath' host port@ within it.
-- 3. Otherwise a TCP connection is made, resolving the host via DNS.
connect :: ByteString -> ByteString -> Int -> IO Transport
connect host hostaddr port
  | not (ByteString.null hostaddr) = connectTcp hostaddr port
  | isUnixSocketHost host = connectUnix (unixSocketPath host port)
  | otherwise = connectTcp host port

-- | Whether a conninfo @host@ value names a Unix-domain socket directory
-- rather than a TCP host - i.e. it looks like an absolute path, per libpq's
-- rule: "If a host name looks like an absolute path name, it specifies
-- Unix-domain communication rather than TCP/IP communication".
isUnixSocketHost :: ByteString -> Bool
isUnixSocketHost host = not (ByteString.null host) && ByteString.head host == 0x2f -- '/'

-- | The path of the socket file libpq expects within a Unix-domain socket
-- directory: @\<directory\>\/.s.PGSQL.\<port\>@.
unixSocketPath :: ByteString -> Int -> FilePath
unixSocketPath directory port = ByteString.Char8.unpack directory <> "/.s.PGSQL." <> show port

-- | Open a TCP connection to the given host and port.
connectTcp :: ByteString -> Int -> IO Transport
connectTcp host port = do
  let hints = Socket.defaultHints {Socket.addrSocketType = Socket.Stream}
  addresses <-
    Socket.getAddrInfo (Just hints) (Just (ByteString.Char8.unpack host)) (Just (show port))
  case addresses of
    [] -> ioError (userError ("Could not resolve host: " <> ByteString.Char8.unpack host))
    address : _ -> do
      sock <- Socket.socket (Socket.addrFamily address) (Socket.addrSocketType address) (Socket.addrProtocol address)
      Socket.connect sock (Socket.addrAddress address)
      buffer <- newReadBuffer
      pure Transport {socket = sock, readBuffer = buffer}

-- | Open a connection to a Unix-domain socket at the given path (the
-- directory\/@\.s\.PGSQL\.\<port\>@ file, per 'unixSocketPath').
connectUnix :: FilePath -> IO Transport
connectUnix path
  | not Socket.isUnixDomainSocketAvailable =
      ioError (userError "pqi-native: Unix-domain sockets are not supported on this platform")
  | otherwise = do
      sock <- Socket.socket Socket.AF_UNIX Socket.Stream 0
      Socket.connect sock (Socket.SockAddrUnix path)
      buffer <- newReadBuffer
      pure Transport {socket = sock, readBuffer = buffer}

-- | An unconnected transport, for representing a \"null\" connection. Its
-- socket is allocated but never connected; it must not be used for I\/O.
unconnected :: IO Transport
unconnected = do
  sock <- Socket.socket Socket.AF_INET Socket.Stream Socket.defaultProtocol
  buffer <- newReadBuffer
  pure Transport {socket = sock, readBuffer = buffer}

-- | Close the connection.
close :: Transport -> IO ()
close transport = Socket.close (socket transport)

-- | The underlying socket file descriptor.
socketFd :: Transport -> IO Int32
socketFd transport = fromIntegral <$> Socket.unsafeFdSocket (socket transport)

-- | Send a serialized message.
--
-- Mirrors libpq's @pqSendSome@ (@fe-misc.c@): a send that cannot complete
-- immediately must not merely block on writability. The server may itself be
-- blocked writing results this connection has not read yet, and a client
-- blocked in @send@ cannot drain them, so both sides would wait forever -
-- the exact pipeline deadlock the comment above libpq's own loop describes.
-- So whenever the socket is not writable, incoming bytes are absorbed into
-- the read buffer first and the wait is for /read-or-write/ readiness:
-- whichever side can make progress does. Sending is preferred whenever the
-- socket is writable (the @orElse@ in 'awaitWritableOrReadable' is
-- left-biased), so a server streaming a large result cannot starve the
-- send.
#if defined(mingw32_HOST_OS)
-- Windows keeps the plain blocking @sendAll@: GHC's I\/O manager has no
-- dependable @threadWaitRead@\/@threadWaitWrite@ there, which the
-- read-while-sending loop is built on. The deadlock the loop prevents is
-- not Windows-specific, but reaching it needs the kernel socket buffers to
-- fill while server output is pending, which the forwarding paths Windows
-- clients typically sit behind make unlikely.
send :: Transport -> Poker.Write -> IO ()
send transport write = Socket.ByteString.sendAll (socket transport) (Poker.toByteString write)
#else
send :: Transport -> Poker.Write -> IO ()
send transport write = go (Poker.toByteString write)
  where
    sock = socket transport
    go bytes
      | ByteString.null bytes = pure ()
      | otherwise =
          awaitWritableOrReadable sock >>= \case
            True -> do
              sent <- Socket.ByteString.send sock bytes
              go (ByteString.drop sent bytes)
            False -> do
              drainIncoming transport
              go bytes

-- | Wait until the socket is writable or, failing that, readable, reporting
-- which side woke ('True' for writable). Both waits are armed at once
-- because a socket that stays unwritable can still keep receiving data the
-- caller must absorb for the server's output path to keep draining.
awaitWritableOrReadable :: Socket.Socket -> IO Bool
awaitWritableOrReadable sock = do
  fd <- Socket.unsafeFdSocket sock
  (writable, cancelWritable) <- threadWaitWriteSTM (fromIntegral fd)
  (readable, cancelReadable) <- threadWaitReadSTM (fromIntegral fd)
  outcome <- atomically (fmap (\_ -> True) writable `orElse` fmap (\_ -> False) readable)
  cancelReadable
  cancelWritable
  pure outcome

-- | Absorb one chunk of incoming data into the read buffer - the analogue
-- of libpq calling @pqReadData()@ inside its send loop. Only called once
-- the socket has reported readable, so the @recv@ cannot block. Room is
-- made for one byte more than is buffered, so the receive always has
-- somewhere to go. The peer closing the connection mid-send is reported
-- like every other transport-level EOF.
drainIncoming :: Transport -> IO ()
drainIncoming transport = do
  buffer <- readIORef (readBuffer transport)
  closed <- receiveSome transport (end buffer - start buffer + 1)
  when closed do
    ioError (mkIOError eofErrorType "pqi-native: connection closed by server" Nothing Nothing)
#endif

-- | Ensure the read buffer holds at least @n@ bytes, pulling from the socket
-- as needed. Throws on EOF before @n@ bytes are available.
--
-- The wait for bytes is deliberately left interruptible. Nothing has been
-- consumed at this point, so a caller that gives up here loses nothing - and
-- is /able/ to give up. A caller blocked on a message the server
-- will never send (an aborted pipeline whose bookkeeping has drifted, say)
-- must stay abandonable by 'System.Timeout.timeout'; masking the wait
-- uninterruptibly turns that stall into a deadlock no timer can break.
--
-- Only the step that moves bytes off the socket and into the buffer is masked
-- (see 'receiveSome'), which is all the atomicity the framing needs, and
-- the blocking wait keeps its abandonability.
fillTo :: Transport -> Int -> IO ()
fillTo transport n = go
  where
    go = do
      buffer <- readIORef (readBuffer transport)
      when (end buffer - start buffer < n) do
        closed <- receiveSome transport n
        when closed do
          ioError (mkIOError eofErrorType "pqi-native: connection closed by server" Nothing Nothing)
        go

-- | Receive one framed message: its type byte and its body (the length prefix,
-- which counts itself, is consumed).
--
-- Buffering the whole frame before consuming any of it keeps the framing
-- atomic without masking the wait: an async exception landing while the frame
-- is still incomplete leaves the buffer untouched, and the frame is taken out
-- of the buffer in a single 'atomicModifyIORef'' step. The body is a slice of
-- the buffer's storage, not a copy. The earlier shape -
-- consuming the header, then blocking again for the body - is what made a
-- mid-read interrupt desync the connection, and what an outer
-- 'Control.Exception.uninterruptibleMask_' was papering over at the cost of
-- making every stall permanent.
receiveFrame :: Transport -> IO (Word8, ByteString)
receiveFrame transport = do
  fillTo transport 5
  header <- ByteString.take 5 . buffered <$> readIORef (readBuffer transport)
  let frameLength = decodeInt32BE (ByteString.drop 1 header)
      frameSize = 5 + max 0 (frameLength - 4)
  fillTo transport frameSize
  atomicModifyIORef' (readBuffer transport) \buffer ->
    let frame = ByteString.take frameSize (buffered buffer)
     in (buffer {start = start buffer + frameSize}, (ByteString.head frame, ByteString.drop 5 frame))

-- | Read bytes until the peer closes the connection (EOF), discarding them.
-- Used by the cancel path to mirror libpq's behaviour: keep the socket open
-- until the server has read the cancel request and closed its end.
readUntilClosed :: Transport -> IO ()
readUntilClosed transport = go
  where
    go = do
      chunk <- Socket.ByteString.recv (socket transport) 4096
      if ByteString.null chunk
        then pure ()
        else go

-- | The numeric IP address of the connected peer (e.g. @\"::1\"@ or
-- @\"127.0.0.1\"@). Throws if the socket has no peer (unconnected).
peerIp :: Transport -> IO ByteString
peerIp transport = do
  addr <- Socket.getPeerName (socket transport)
  (Just ip, _) <- Socket.getNameInfo [Socket.NI_NUMERICHOST] True False addr
  pure (ByteString.Char8.pack ip)

decodeInt32BE :: ByteString -> Int
decodeInt32BE = ByteString.foldl' (\acc w -> acc * 256 + fromIntegral w) 0 . ByteString.take 4