Своя СУБД на Haskell · урок 11: бинарный протокол и сервер

Кому это. За десять уроков мы построили на Haskell целый движок: страницы, буферный пул, кучу, B+Tree, парсер SQL, планировщик и исполнитель, журнал, транзакции, блокировки и vacuum. В прошлом уроке он научился убирать за собой. Но говорить с движком умеет только тот, кто импортировал OurDb.Engine.Database в свою программу. Снаружи к нашей СУБД не подключиться.

Зачем именно Haskell. Сегодня выставляем движок в сеть: свой бинарный протокол в духе Postgres, TCP-сервер и клиент. На Haskell урок держится на одном решении: весь протокол это обработчик байт. Кадрирование потока это чистая функция «состояние и чанк в кадры и новое состояние». Кодеки сообщений чистые. Сессия соединения превращает SQL-строку в список сообщений поверх Database. А соединение целиком это функция «чанк байт в ответ и флаг закрытия». Сокетов в основной библиотеке нет вовсе. Пакет network нужен только в отдельной подбиблиотеке ourdb:net, где сервер читает чанки из сокета, зовёт обработчик и пишет ответ.

Важно про объём и про этап. Урок аддитивный: движок не меняется, сверху появляются модули OurDb.Wire.* и подбиблиотека ourdb:net. И одно практическое следствие решения выше. Раннер курса, где проверяются этапы, собирает код только из boot-пакетов GHC, а network среди них нет. Поэтому финальный этап проверяет обработчик соединения над байтами: тесты кормят его кадрами, нарезанными как угодно, и разбирают ответ. Сетевую обвязку из ourdb:net мы разберём как материал для чтения, а запустить её можно локально в эталоне через cabal test.

Где мы в движке

┌────────────────────────────────────────────────┐
│  Бинарный протокол + TCP-сервер                  │  урок 11  ← мы здесь
├────────────────────────────────────────────────┤
│  Исполнитель (Volcano) + EXPLAIN                 │  урок 6
│  Планировщик, binder, каталог                    │  урок 6
│  Парсер SQL                                      │  урок 5
├────────────────────────────────────────────────┤
│  Vacuum (сборка мёртвых версий)                  │  урок 10
│  Блокировки (2PL), конкуренция                   │  урок 9
│  MVCC (видимость, версии кортежа)                │  урок 8
│  Менеджер транзакций (xid, состояния, undo)      │  урок 8
│  WAL (журнал, redo, recovery)                    │  урок 7
├────────────────────────────────────────────────┤
│  B+Tree индекс                                   │  урок 4
│  Куча (slotted pages, seq scan)                  │  урок 3
├────────────────────────────────────────────────┤
│  Буферный пул (кэш страниц, вытеснение)           │  урок 2
│  DiskManager (страницы поверх файла)              │  урок 1
└────────────────────────────────────────────────┘

Верхний слой лежит над движком, а не внутри него. Раскладка модулей такая:

src/OurDb/Wire/Framing.hs    кадры из потока байт (чистое состояние)
src/OurDb/Wire/Messages.hs   кодеки сообщений (чистые)
src/OurDb/Wire/Session.hs    SQL в серверные сообщения, транзакция соединения
src/OurDb/Wire/Conn.hs       соединение сервера над байтами
src/OurDb/Wire/Client.hs     клиентский аккумулятор ответа (чистый)
net/OurDb/Net/Server.hs      ourdb:net, TCP-сервер, поток на соединение
net/OurDb/Net/Client.hs      ourdb:net, TCP-клиент

Идея протокола

Мы берём simple query protocol Postgres, урезанный до сути. Клиент шлёт два вида сообщений, сервер пять.

Направление	Тип	Байт	Смысл
клиент	`Query`	`Q`	текст SQL с завершающим NUL
клиент	`Terminate`	`X`	закрыть соединение
сервер	`RowDescription`	`T`	имена и типы колонок
сервер	`DataRow`	`D`	одна строка, значения текстом
сервер	`CommandComplete`	`C`	тег: `SELECT 2`, `INSERT 0 1`
сервер	`ErrorResponse`	`E`	код в духе SQLSTATE и текст
сервер	`ReadyForQuery`	`Z`	готов к запросу, статус транзакции

Каждый ответ заканчивается ReadyForQuery с одним байтом статуса: I (вне транзакции), T (в транзакции), E (транзакция упала, до ROLLBACK ничего нельзя). По этому байту клиент понимает, что ответ закончился, и в каком состоянии сессия.

Кадрирование

TCP это поток байт, а не поток сообщений. Один recv может вернуть половину сообщения, три сообщения сразу или полтора. Поэтому каждое сообщение упаковано в кадр: байт типа, четыре байта длины big-endian и payload. Длина включает свои 4 байта, но не байт типа, как в Postgres.

-- src/OurDb/Wire/Framing.hs

-- | Кадрирование протокола (урок 11): u8 тип, u32 длина big-endian (включает 4 байта
-- длины, не включает байт типа), payload длиной length - 4.
--
-- 'FrameReader' превращает поток произвольно нарезанных TCP-чанков в целые кадры: один
-- чанк может нести несколько кадров или половину кадра. Это чистое состояние: push
-- возвращает собравшиеся кадры и новый остаток.
module OurDb.Wire.Framing
  ( Frame (..)
  , encodeFrame
  , FrameReader
  , newFrameReader
  , pushChunk
  , pendingBytes
  ) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as B
import Data.Word (Word8)
import OurDb.Bytes (build, getU32BE, runGetAt)

data Frame = Frame {frameType :: !Word8, framePayload :: !BS.ByteString}
  deriving (Eq, Show)

frameHeaderSize :: Int
frameHeaderSize = 5

encodeFrame :: Word8 -> BS.ByteString -> BS.ByteString
encodeFrame kind payload =
  build (B.word8 kind <> B.word32BE (fromIntegral (4 + BS.length payload)) <> B.byteString payload)

newtype FrameReader = FrameReader BS.ByteString

newFrameReader :: FrameReader
newFrameReader = FrameReader BS.empty

-- | Недочитанный хвост (для тестов и проверки "ровно один кадр").
pendingBytes :: FrameReader -> BS.ByteString
pendingBytes (FrameReader buffer) = buffer

-- | Добавить чанк и вернуть все кадры, собравшиеся целиком.
pushChunk :: FrameReader -> BS.ByteString -> ([Frame], FrameReader)
pushChunk (FrameReader buffer) chunk = go (buffer <> chunk) []
  where
    go bytes acc
      | BS.length bytes < frameHeaderSize = (reverse acc, FrameReader bytes)
      | otherwise =
          let total = either (const 0) ((+ 1) . fromIntegral) (runGetAt getU32BE bytes 1)
           in if BS.length bytes < total
                then (reverse acc, FrameReader bytes) -- кадр ещё не дочитан
                else
                  let frame = Frame (BS.head bytes) (BS.drop frameHeaderSize (BS.take total bytes))
                   in go (BS.drop (max 1 total) bytes) (frame : acc)

В Effect- и TS-версиях этого урока FrameReader это изменяемый объект с буфером внутри. Здесь он значение: newtype над недочитанным хвостом. pushChunk берёт старое состояние и чанк, а возвращает пару: все целые кадры и новое состояние. Та же форма s -> (a, s), что у менеджера блокировок в уроке 9. Кто хранит состояние между чанками, решает вызывающий: сервер положит его в IORef, клиент в MVar, а тест просто протянет значения руками.

go отрезает кадры, пока хватает байт. Меньше пяти байт: заголовок не дочитан, ждём. Длина есть, но байт меньше total: кадр не дочитан, ждём. Иначе отрезаем кадр и идём дальше. total это длина из заголовка плюс байт типа. max 1 total гарантирует продвижение даже на мусорном заголовке с нулевой длиной: цикл не застрянет на одном месте. BS.take и BS.drop не копируют байты, поэтому кадры это срезы одного буфера.

Сообщения

Кодеки сообщений тоже чистые. Самое интересное в них это тип серверного сообщения.

-- src/OurDb/Wire/Messages.hs

-- | Сообщения протокола (урок 11): simple query protocol в духе Postgres, числа
-- big-endian. Чистые кодеки payload-ов; кадрирование живёт в "OurDb.Wire.Framing".
--
-- Значения в DataRow это текст: int десятичной строкой, bool как t/f, text как есть.
-- Серверное сообщение параметризовано типом строки: сервер шлёт 'Value', клиент получает
-- сырой текст (Maybe Text) и разбирает его по типам из RowDescription.

-- | Статус транзакции в ReadyForQuery.
data TxnStatus = StatusIdle | StatusInTx | StatusFailed
  deriving (Eq, Show)

statusByte :: TxnStatus -> Word8
statusByte = \case
  StatusIdle -> 0x49 -- 'I'
  StatusInTx -> 0x54 -- 'T'
  StatusFailed -> 0x45 -- 'E'

msgQuery, msgTerminate, msgRowDescription, msgDataRow, msgCommandComplete, msgErrorResponse, msgReadyForQuery :: Word8
msgQuery = 0x51 -- 'Q'
msgTerminate = 0x58 -- 'X'
msgRowDescription = 0x54 -- 'T'
msgDataRow = 0x44 -- 'D'
msgCommandComplete = 0x43 -- 'C'
msgErrorResponse = 0x45 -- 'E'
msgReadyForQuery = 0x5a -- 'Z'

data ClientMessage = Query !Text | Terminate
  deriving (Eq, Show)

data ServerMsg row
  = RowDescription ![Column]
  | DataRow !row
  | CommandComplete !Text
  | -- | Код (SQLSTATE-подобный) и текст.
    ErrorResponse !Text !Text
  | ReadyForQuery !TxnStatus
  deriving (Eq, Show)

type ServerMessage = ServerMsg [Value]

type DecodedServerMessage = ServerMsg [Maybe Text]

ServerMsg параметризован типом строки. Сервер собирает ServerMessage, где DataRow несёт типизированные [Value] прямо из исполнителя. На проводе значения превращаются в текст. Клиент разбирает кадр в DecodedServerMessage, где та же строка это [Maybe Text]: сырой текст ячеек, Nothing для NULL. Разбирать текст по типам клиент будет сам, по RowDescription. Один тип данных описывает обе стороны, а параметр говорит, на какой стороне провода ты стоишь. Перепутать их нельзя: функция, ждущая ServerMessage, не примет DecodedServerMessage.

-- src/OurDb/Wire/Messages.hs (продолжение)

clientMessageType :: ClientMessage -> Word8
clientMessageType = \case
  Query _ -> msgQuery
  Terminate -> msgTerminate

-- | SQL плюс завершающий NUL.
encodeClientPayload :: ClientMessage -> BS.ByteString
encodeClientPayload = \case
  Query sql -> utf8 sql <> BS.singleton 0
  Terminate -> BS.empty

decodeClientMessage :: Word8 -> BS.ByteString -> Either Text ClientMessage
decodeClientMessage kind payload
  | kind == msgQuery =
      let body = if not (BS.null payload) && BS.last payload == 0 then BS.init payload else payload
       in Right (Query (fromUtf8 body))
  | kind == msgTerminate = Right Terminate
  | otherwise = Left ("неизвестный тип клиентского сообщения 0x" <> T.pack (showHex kind ""))

serverMessageType :: ServerMsg row -> Word8
serverMessageType = \case
  RowDescription _ -> msgRowDescription
  DataRow _ -> msgDataRow
  CommandComplete _ -> msgCommandComplete
  ErrorResponse _ _ -> msgErrorResponse
  ReadyForQuery _ -> msgReadyForQuery

encodeServerPayload :: ServerMessage -> BS.ByteString
encodeServerPayload = \case
  RowDescription fields -> build (u16 (length fields) <> foldMap field fields)
  DataRow values -> build (u16 (length values) <> foldMap cell values)
  CommandComplete tag -> utf8 tag
  ErrorResponse code message -> let c = utf8 code in build (u16 (BS.length c) <> B.byteString c <> B.byteString (utf8 message))
  ReadyForQuery status -> BS.singleton (statusByte status)
  where
    u16 = B.word16BE . fromIntegral
    field (Column name t) = let n = utf8 name in u16 (BS.length n) <> B.byteString n <> B.word8 (typeCode t)
    cell value = case encodeValueText value of
      Nothing -> B.int32BE (-1)
      Just bytes -> B.int32BE (fromIntegral (BS.length bytes)) <> B.byteString bytes

serverMessageType работает для любого row: тип кадра не зависит от того, что лежит в строке. А encodeServerPayload принимает только серверную сторону, ServerMessage. Ячейка DataRow это длина i32 и байты, а длина -1 означает NULL, как в Postgres. Функции field и cell возвращают Builder, и foldMap склеивает их в один кусок без промежуточных строк.

Разбор серверного сообщения нужен клиенту и тестам:

-- src/OurDb/Wire/Messages.hs (продолжение)

-- | Разобрать серверное сообщение (DataRow остаётся сырым текстом).
decodeServerMessage :: Word8 -> BS.ByteString -> Either Text DecodedServerMessage
decodeServerMessage kind payload
  | kind == msgRowDescription = runGet (RowDescription <$> counted field) payload
  | kind == msgDataRow = runGet (DataRow <$> counted cell) payload
  | kind == msgCommandComplete = Right (CommandComplete (fromUtf8 payload))
  | kind == msgErrorResponse = runGet errorResponse payload
  | kind == msgReadyForQuery = runGet (getU8 >>= either fail' (pure . ReadyForQuery) . status) payload
  | otherwise = Left ("неизвестный тип серверного сообщения 0x" <> T.pack (showHex kind ""))
  where
    fail' = fail . T.unpack
    counted item = getU16BE >>= \n -> replicateM (fromIntegral n) item
    field = do
      name <- getU16BE >>= fmap fromUtf8 . getBytes . fromIntegral
      code <- getU8
      case code of
        0 -> pure (Column name TInt)
        1 -> pure (Column name TText)
        2 -> pure (Column name TBool)
        _ -> fail ("неизвестный код типа колонки " <> show code)
    cell =
      getI32BE >>= \case
        -1 -> pure Nothing
        len -> Just . fromUtf8 <$> getBytes (fromIntegral len)
    errorResponse = do
      code <- getU16BE >>= fmap TE.decodeLatin1 . getBytes . fromIntegral
      ErrorResponse code . fromUtf8 <$> getRest
    status byte = case chr (fromIntegral byte) of
      'I' -> Right StatusIdle
      'T' -> Right StatusInTx
      'E' -> Right StatusFailed
      c -> Left ("неизвестный статус транзакции \"" <> T.singleton c <> "\"")

-- | Текстовое представление значения для DataRow (Nothing это NULL).
encodeValueText :: Value -> Maybe BS.ByteString
encodeValueText = \case
  VNull -> Nothing
  VInt n -> Just (utf8 (T.pack (show n)))
  VBool b -> Just (if b then "t" else "f")
  VText s -> Just (utf8 s)

-- | Разобрать текстовое значение колонки по её типу.
decodeValueText :: ColumnType -> Text -> Either Text Value
decodeValueText columnType text = case columnType of
  TText -> Right (VText text)
  TBool -> Right (VBool (text == "t"))
  TInt -> case TR.signed TR.decimal text of
    Right (n, rest) | T.null rest, n >= toInteger (minBound :: Int64) && n <= toInteger (maxBound :: Int64) -> Right (VInt (fromInteger n))
    _ -> Left ("не int: " <> text)

counted item это помощник «прочитай счётчик u16, потом столько же элементов»: так устроены и RowDescription, и DataRow. Ошибка разбора внутри Get (неизвестный код колонки, обрыв payload) превращается в Left, а не в падение процесса: мусорный кадр от сервера станет WireProtocolError у клиента. decodeValueText для INT разбирает число в Integer и только потом проверяет диапазон Int64. Иначе fromInteger молча обрезал бы 99999999999999999999 по модулю 2^64.

Сессия соединения

Сессия превращает строку SQL в список серверных сообщений и помнит транзакцию своего соединения.

-- src/OurDb/Wire/Session.hs

-- | Сессия соединения (урок 11): SQL-строка в последовательность серверных сообщений,
-- состояние одной транзакции соединения и статус ReadyForQuery. Без сокетов: сервер
-- вызывает 'runSession' на каждый Query-кадр.
--
-- Каждое соединение это своя транзакция через 'begin' движка, а не общий sessionXid
-- (иначе соединения делили бы одну транзакцию). Упавшая транзакция (ошибка внутри BEGIN)
-- переходит в failed: до ROLLBACK любая команда отвергается, как в Postgres.

data Session = Session
  { sessDb :: !Database
  , sessHandle :: !(IORef (Maybe Transaction))
  , sessFailed :: !(IORef Bool)
  }

newSession :: Database -> IO Session
newSession db = Session db <$> newIORef Nothing <*> newIORef False

-- | Выполнить SQL и вернуть ответ, завершённый ReadyForQuery.
runSession :: Session -> Text -> IO [ServerMessage]
runSession session sql = case parseCommand sql of
  Left err -> respond session [errorOf (Right err)]
  Right (CmdTxn op) -> runTxn session op
  Right _ -> do
    failed <- readIORef (sessFailed session)
    if failed
      then respond session [ErrorResponse "25P02" "текущая транзакция прервана"]
      else do
        handle <- readIORef (sessHandle session)
        outcome <- trySync (maybe (execute (sessDb session) sql) (`txExecute` sql) handle)
        case outcome of
          Right result -> respond session (messagesFor result)
          Left err -> do
            mapM_ (const (writeIORef (sessFailed session) True)) handle
            respond session [errorOf err]

Помнишь dbSessionXid из урока 8? Это одна сессионная транзакция на весь движок. Для сервера она не годится: два соединения делили бы одну транзакцию. Поэтому у сессии свой хэндл Transaction из урока 9, полученный через begin. Нет хэндла, и оператор идёт через execute в autocommit. Есть, и через txExecute в транзакции соединения. Выбирает maybe: «если Nothing, первое, иначе второе применить к хэндлу».

Синтаксическая ошибка ловится ещё до движка: parseCommand вернул Left, ответ это ErrorResponse и ReadyForQuery с прежним статусом. Ошибка внутри транзакции переводит сессию в состояние failed: mapM_ (const ...) handle это опять идиома «сделай, если есть», флаг поднимается только при открытом хэндле. Ошибка autocommit-оператора транзакцию не роняет, её и нет.

Самое тонкое место сессии это trySync:

-- src/OurDb/Wire/Session.hs (продолжение)

-- | Синхронные исключения как значение, асинхронные (killThread) летят дальше.
trySync :: IO a -> IO (Either (Either SomeException DbError) a)
trySync action =
  try action >>= \case
    Right value -> pure (Right value)
    Left e
      | Just (_ :: SomeAsyncException) <- fromException e -> throwIO e
      | Just db <- fromException e -> pure (Left (Right db))
      | otherwise -> pure (Left (Left e))

errorOf :: Either SomeException DbError -> ServerMessage
errorOf = \case
  Right err -> ErrorResponse (errorCode err) (errorMessage err)
  Left other -> ErrorResponse "XX000" (T.pack (displayException other))

-- | SQLSTATE-подобный код по виду ошибки движка.
errorCode :: DbError -> Text
errorCode = \case
  UniquenessError _ -> "23505"
  SerializationError _ -> "40001"
  DeadlockError _ -> "40P01"
  LockWaitError _ -> "55P03"
  SqlError _ _ -> "42601"
  EngineError _ -> "42000"
  StorageError _ -> "XX000"

Сервер должен ответить клиенту на любую ошибку оператора, а не уронить соединение. Отсюда try с SomeException. Но ловить всё подряд опасно. В Haskell есть асинхронные исключения: killThread бросает исключение в чужой поток, и так сервер при остановке гасит потоки соединений. Если trySync проглотит такое исключение и превратит его в ErrorResponse, поток соединения выживет после killThread, и сервер никогда не дождётся его завершения. Поэтому первым проверяется SomeAsyncException, и его throwIO бросает дальше. DbError становится типизированной ошибкой с SQLSTATE-кодом, всё остальное (например, IOError от упавшего fsync) идёт как XX000. Это прямое продолжение урока про исключения и ресурсы.

errorCode это полное сопоставление по DbError. Добавишь в движок новый вид ошибки, и -Wall укажет на это место: протокол не забудет дать ей код.

Транзакционные команды сессия обрабатывает сама:

-- src/OurDb/Wire/Session.hs (продолжение)

-- | Откатить открытую транзакцию (Terminate и закрытие сокета).
closeSession :: Session -> IO ()
closeSession session =
  readIORef (sessHandle session) >>= \case
    Nothing -> pure ()
    Just handle -> do
      writeIORef (sessHandle session) Nothing
      writeIORef (sessFailed session) False
      txRollback handle

runTxn :: Session -> TxnOp -> IO [ServerMessage]
runTxn session = \case
  TxnBegin -> do
    readIORef (sessHandle session) >>= \case
      Just _ -> pure ()
      Nothing -> do
        handle <- begin (sessDb session)
        writeIORef (sessHandle session) (Just handle)
        writeIORef (sessFailed session) False
    respond session [CommandComplete "BEGIN"]
  TxnCommit -> finishTxn True >>= respond session . pure . CommandComplete
  TxnRollback -> finishTxn False >> respond session [CommandComplete "ROLLBACK"]
  where
    -- COMMIT упавшей транзакции это откат с тегом ROLLBACK
    finishTxn :: Bool -> IO Text
    finishTxn wantCommit =
      readIORef (sessHandle session) >>= \case
        Nothing -> pure (if wantCommit then "COMMIT" else "ROLLBACK")
        Just handle -> do
          writeIORef (sessHandle session) Nothing
          wasFailed <- atomicModifyIORef' (sessFailed session) (False,)
          if not wantCommit || wasFailed
            then "ROLLBACK" <$ txRollback handle
            else "COMMIT" <$ txCommit handle

respond :: Session -> [ServerMessage] -> IO [ServerMessage]
respond session messages = (\status -> messages <> [ReadyForQuery status]) <$> sessionStatus session

sessionStatus :: Session -> IO TxnStatus
sessionStatus session = do
  failed <- readIORef (sessFailed session)
  handle <- readIORef (sessHandle session)
  pure $ case (failed, handle) of
    (True, _) -> StatusFailed
    (_, Nothing) -> StatusIdle
    _ -> StatusInTx

-- | Развернуть результат в серверные сообщения (без ReadyForQuery).
messagesFor :: ExecResult -> [ServerMessage]
messagesFor = \case
  ResSelect _ fields rows -> RowDescription fields : map DataRow rows <> [CommandComplete ("SELECT " <> T.pack (show (length rows)))]
  ResCommand tag -> [CommandComplete tag]
  ResExplain plan -> [RowDescription [Column "QUERY PLAN" TText], DataRow [VText plan], CommandComplete "EXPLAIN"]

Повторный BEGIN при открытой транзакции молча проходит, как предупреждение в Postgres. COMMIT упавшей транзакции откатывает её и отвечает тегом ROLLBACK: клиент видит по тегу, что зафиксировать не вышло. atomicModifyIORef' (sessFailed session) (False,) за один шаг читает старый флаг и сбрасывает его: секция (False,) это функция \old -> (False, old), она нужна расширению TupleSections, которое входит в GHC2021. respond дописывает ReadyForQuery к любому ответу, поэтому ни одна ветка не забудет его отправить. EXPLAIN уходит одной текстовой строкой в колонке QUERY PLAN, как в psql.

дальше после входа

Урок целиком открыт всем, кто вошёл. Войди по своему email, чтобы читать дальше.

Вход пока по приглашению. Если хочешь учиться, напиши Vassiliy.

войти