Stale.hs

-- | 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 = 30

spec :: Pqi.Adapter -> SpecWith ByteString
spec adapter =
  describe "stale cancel" do
    it "does not corrupt the next command when a cancel is sent during pipeline clean-up" \conninfo -> do
      outcomes <- for [1 .. iterations] \i -> do
        outcome <- runScenario adapter 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 = [(i, outcome) | (i, Just outcome) <- 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 conninfo =
  bracket (Lq.connectdb adapter conninfo) Lq.finish \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.
    _ <- Lq.enterPipelineMode 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
      Nothing -> pure (Just (Lq.FatalError, Just "no-result"))
      Just observation ->
        case status observation of
          Lq.TuplesOk -> pure Nothing
          status -> pure (Just (status, fromMaybe Nothing (lookup Lq.DiagSqlstate (errorFields observation))))