-- | Deterministic regression coverage for the stale-cancel bug: a cancel
-- dispatched during pipeline clean-up must not corrupt the next command run on
-- the connection.
--
-- Root cause: 'Pqi.cancel' opens a TCP connection to the postmaster, sends the
-- @CancelRequest@, and — in the buggy implementation — immediately closes its
-- end.  The request can then sit unread in the postmaster's socket buffer while
-- the client races ahead and issues its next query.  The postmaster eventually
-- reads the request and delivers @SIGINT@ to the backend, which by then is
-- executing the /next/ statement, cancelling it with SQLSTATE @57014@.
--
-- The race is reliably opened by the pipeline clean-up sequence that
-- @hasql@ runs after a timeout fires mid-read: once the pipeline is aborted,
-- 'Pqi.getResult' returns synthetic results for the outstanding commands
-- __without blocking on the wire__ (see @getNextResult@ in
-- @Pqi.Native.Query@).  That lets the client reach its follow-up query before
-- the postmaster has processed the cancel.
--
-- The fix mirrors libpq's @PQcancel@: after sending the request, drain the
-- cancel socket until the server closes its end (which it only does after the
-- signal has been dispatched), so the cancel can no longer land on a later
-- command.
--
-- A single run hits the race only intermittently (the postmaster usually
-- processes the cancel in time), so this spec replays the scenario many times.
-- On the buggy implementation at least one iteration reliably loses the race
-- and reports @57014@ on the victim query; the fixed implementation passes
-- every iteration.
module Pqi.Conformance.Operation.Cancel.Stale
  ( spec,
  )
where

import Control.Exception (bracket)
import qualified Pqi
import qualified Pqi as Lq
import Pqi.Conformance.Observation (ResultObservation (..))
import Pqi.Conformance.Prelude
import Pqi.Conformance.Scenario (drainResults, execScenario, float8Oid)
import System.Timeout (timeout)
import Test.Hspec

-- | How many times to replay the racy scenario.  Each iteration that loses the
-- race on buggy code surfaces SQLSTATE @57014@ on the victim query; enough
-- iterations make at least one failure overwhelmingly likely (the underlying
-- per-run failure rate on buggy code is roughly 30%).
iterations :: Int
iterations :: Int
iterations = Int
30

spec :: Pqi.Adapter -> SpecWith ByteString
spec :: Adapter -> SpecWith ByteString
spec Adapter
adapter =
  String -> SpecWith ByteString -> SpecWith ByteString
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"stale cancel" do
    String
-> (ByteString -> IO ()) -> SpecWith (Arg (ByteString -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"does not corrupt the next command when a cancel is sent during pipeline clean-up" \ByteString
conninfo -> do
      outcomes <- [Int]
-> (Int -> IO (Int, Maybe (ExecStatus, Maybe ByteString)))
-> IO [(Int, Maybe (ExecStatus, Maybe ByteString))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Int
1 .. Int
iterations] \Int
i -> do
        outcome <- Adapter -> ByteString -> IO (Maybe (ExecStatus, Maybe ByteString))
runScenario Adapter
adapter ByteString
conninfo
        pure (i, outcome)
      -- Every iteration's victim query must have succeeded.  Any iteration that
      -- did not means a stale cancel corrupted the connection (57014).
      let corrupted = [(Int
i, (ExecStatus, Maybe ByteString)
outcome) | (Int
i, Just (ExecStatus, Maybe ByteString)
outcome) <- [(Int, Maybe (ExecStatus, Maybe ByteString))]
outcomes]
      corrupted `shouldBe` []

-- | Run one stale-cancel scenario on a fresh connection.  Returns 'Nothing' if
-- the victim query succeeded, or @Just (status, sqlstate)@ describing how it
-- failed (e.g. @(FatalError, Just "57014")@).
runScenario ::
  Pqi.Adapter ->
  ByteString ->
  IO (Maybe (Lq.ExecStatus, Maybe ByteString))
runScenario :: Adapter -> ByteString -> IO (Maybe (ExecStatus, Maybe ByteString))
runScenario Adapter
adapter ByteString
conninfo =
  IO Connection
-> (Connection -> IO ())
-> (Connection -> IO (Maybe (ExecStatus, Maybe ByteString)))
-> IO (Maybe (ExecStatus, Maybe ByteString))
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Adapter -> ByteString -> IO Connection
Lq.connectdb Adapter
adapter ByteString
conninfo) Connection -> IO ()
Lq.finish \Connection
connection -> do
    -- Build a pipeline with a fast query followed by a slow one, mirroring the
    -- sequence hasql issues.  Two prepared statements keep pendingParses in
    -- play, matching the real-world reproduction.
    _ <- Connection -> IO Bool
Lq.enterPipelineMode Connection
connection
    _ <- Lq.sendPrepare connection "s1" "select $1::int" Nothing
    _ <- Lq.sendQueryPrepared connection "s1" [Just ("42", Lq.Text)] Lq.Text
    _ <- Lq.sendPrepare connection "s2" "select pg_sleep($1)" (Just [float8Oid])
    _ <- Lq.sendQueryPrepared connection "s2" [Just ("0.1", Lq.Text)] Lq.Text
    _ <- Lq.pipelineSync connection

    -- Interrupt the read mid-pipeline, exactly as hasql's timeout does: the
    -- slow pg_sleep is still running when the read is abandoned.
    _ <- timeout 50_000 (drainResults connection)

    -- cleanUpAfterInterruption: drain, cancel, drain.  The cancel is sent while
    -- the pipeline is mid-flight; on buggy code its SIGINT can arrive late.
    _ <- drainResults connection
    mHandle <- Lq.getCancel connection
    _ <- for mHandle Lq.cancel
    _ <- drainResults connection

    -- leavePipeline: the exact restore sequence hasql performs, including the
    -- retry it falls back to.
    statusBefore <- Lq.pipelineStatus connection
    when (statusBefore == Lq.PipelineOn) do
      _ <- Lq.pipelineSync connection
      _ <- drainResults connection
      _ <- Lq.sendFlushRequest connection
      _ <- drainResults connection
      ok <- Lq.exitPipelineMode connection
      unless ok do
        _ <- drainResults connection
        void (Lq.exitPipelineMode connection)

    -- The victim.  It must not be cancelled by the stale signal.
    execScenario "select 99" connection >>= \case
      Maybe ResultObservation
Nothing -> Maybe (ExecStatus, Maybe ByteString)
-> IO (Maybe (ExecStatus, Maybe ByteString))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((ExecStatus, Maybe ByteString)
-> Maybe (ExecStatus, Maybe ByteString)
forall a. a -> Maybe a
Just (ExecStatus
Lq.FatalError, ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"no-result"))
      Just ResultObservation
observation ->
        case ResultObservation -> ExecStatus
status ResultObservation
observation of
          ExecStatus
Lq.TuplesOk -> Maybe (ExecStatus, Maybe ByteString)
-> IO (Maybe (ExecStatus, Maybe ByteString))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (ExecStatus, Maybe ByteString)
forall a. Maybe a
Nothing
          ExecStatus
status -> Maybe (ExecStatus, Maybe ByteString)
-> IO (Maybe (ExecStatus, Maybe ByteString))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((ExecStatus, Maybe ByteString)
-> Maybe (ExecStatus, Maybe ByteString)
forall a. a -> Maybe a
Just (ExecStatus
status, Maybe ByteString -> Maybe (Maybe ByteString) -> Maybe ByteString
forall a. a -> Maybe a -> a
fromMaybe Maybe ByteString
forall a. Maybe a
Nothing (FieldCode
-> [(FieldCode, Maybe ByteString)] -> Maybe (Maybe ByteString)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup FieldCode
Lq.DiagSqlstate (ResultObservation -> [(FieldCode, Maybe ByteString)]
errorFields ResultObservation
observation))))