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

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

Зачем именно Go. Это финал воркшопа: выставляем движок в сеть. Поднимем TCP-сервер, придумаем своё кадрирование и реализуем simple query protocol в духе Postgres. Для Go это родная задача: сетевой сервер на горутине-на-соединение это то, ради чего язык делали. Каждое соединение читает свой сокет блокирующим Read в своей горутине, а рантайм под капотом мультиплексирует их через epoll. И тут наконец оживает всё, что мы готовили: движок встаёт под sync.RWMutex, мьютексы пула и журнала, ждавшие с уроков 2 и 7, встречают настоящих конкурентов, а go test -race получает работу.

Важно про объём. Урок почти целиком аддитивный: новый пакет wire из пяти файлов и одна правка движка, sync.RWMutex вокруг engine.Database. Кадрирование, кодеки сообщений, сессия, сервер и клиент. Сервер опирается на готовое: db.Execute, хэндлы db.Begin из урока 8, локи строк из урока 9 и коды ошибок из dberr. Главные новые навыки тут не про базы данных, а про сеть и конкурентность Go: сборка кадров из рваного потока, big-endian, состояние сессии, жизненный цикл горутин и их честное завершение.

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

wire/              протокол и TCP-сервер                   урок 11  ← мы здесь
engine/            каталог, binder, планировщик, Volcano   урок 6, +RWMutex сегодня
sql/               лексер, парсер, AST                     урок 5
vacuum/            сборка мёртвых версий                   урок 10
lock/              блокировки, граф ожидания               урок 9
txn/               xid, снимки, undo-log                   урок 8
dberr/             ошибки движка с кодом                   урок 8
wal/, recovery/    журнал и redo-проход                    урок 7
btree/             B+Tree индекс                           урок 4
heap/, tuple/      куча и кортеж                           урок 3
bufferpool/        кэш страниц, вытеснение clock           урок 2
page/, disk/       страница и файл                         урок 1

Протокол не лезет внутрь движка: он сидит сверху и переводит байты сети в вызовы engine.Database и обратно. Каждое соединение получает свою горутину, свою сессию и свою транзакцию. Изоляцию между соединениями дают снимки из урока 9, а целостность общих структур даст RWMutex движка, который мы поставим сегодня.

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

Простой протокол запросов это диалог сообщений с типами. Все числа в нём big-endian. Это сознательный контраст с диском: на странице из урока 1 мы писали binary.LittleEndian, а на проводе пишем binary.BigEndian. Пакет encoding/binary даёт оба порядка одним интерфейсом, меняется одна переменная.

Диалог строго запрос-ответ. Клиент шлёт Query с SQL-строкой. Сервер отвечает последовательностью: для SELECT это RowDescription (имена и типы колонок), потом DataRow на каждую строку, потом CommandComplete с тегом (SELECT 3). Любую последовательность закрывает ReadyForQuery со статусом транзакции: I вне транзакции, T внутри, E транзакция упала. По нему клиент понимает, что ответ дочитан. Ошибка это ErrorResponse с кодом и текстом, а за ним всё равно ReadyForQuery.

Значения в DataRow идут текстом: int десятичной строкой, bool как t или f, text как есть, NULL это длина минус один. Клиент разбирает текст обратно по типу колонки из RowDescription. Это упрощение Postgres, но его хватает для полного цикла.

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

Начнём снизу. Кадр это байт типа (ASCII), затем u32 длины (big-endian, включает свои 4 байта, но не байт типа), затем payload. Заведи пакет wire и файл wire/framing.go.

// Package wire это бинарный протокол и сервер (урок 11): simple query protocol в духе
// Postgres, числа big-endian (SPEC.md, раздел "Протокол"). Кадр это u8 тип, u32 длина
// (включает 4 байта длины, не включает байт типа) и payload длиной length - 4.
package wire

import (
	"encoding/binary"
	"fmt"
)

const frameHeaderSize = 5

var be = binary.BigEndian

// Frame это один разобранный кадр.
type Frame struct {
	Type    byte
	Payload []byte
}

// EncodeFrame собирает кадр: байт типа, длина (4 плюс payload), payload.
func EncodeFrame(typ byte, payload []byte) []byte {
	out := make([]byte, frameHeaderSize, frameHeaderSize+len(payload))
	out[0] = typ
	be.PutUint32(out[1:], uint32(4+len(payload)))
	return append(out, payload...)
}

// FrameReader превращает поток произвольно нарезанных TCP-чанков в целые кадры: один чанк
// может нести несколько кадров или половину кадра.
type FrameReader struct {
	buf []byte
}

// Push добавляет чанк и возвращает все собравшиеся кадры. Длина меньше 4 это поломанный
// поток, а не "кадр ещё не дочитан".
func (r *FrameReader) Push(chunk []byte) ([]Frame, error) {
	r.buf = append(r.buf, chunk...)
	var frames []Frame
	for len(r.buf) >= frameHeaderSize {
		length := be.Uint32(r.buf[1:])
		if length < 4 {
			return frames, fmt.Errorf("кадр 0x%x с длиной %d меньше 4", r.buf[0], length)
		}
		// shortcut: длина не ограничена сверху, чужой клиент может заставить копить память
		total := 1 + int(length)
		if len(r.buf) < total {
			break
		}
		frames = append(frames, Frame{Type: r.buf[0], Payload: r.buf[frameHeaderSize:total:total]})
		r.buf = r.buf[total:]
	}
	return frames, nil
}

FrameReader это самая важная часть урока и та, про которую легко забыть. TCP даёт поток байт, а не сообщения: один conn.Read может вернуть половину кадра, полтора кадра или три сразу. Push копит байты и отрезает ровно те кадры, что собрались целиком, оставляя хвост на потом. Без этого приёмник рано или поздно прочитает мусор: на медленной сети или большом ответе фрагментация неизбежна. Длина меньше 4 это не “кадр ещё не дочитан”, а сломанный поток: payload отрицательной длины не бывает, и ждать дальше бессмысленно.

Посмотри на r.buf[frameHeaderSize:total:total]. Это полное выражение среза с третьим индексом: оно ограничивает ёмкость payload ровно его длиной. Payload это окно в общий буфер чтения, и сразу за ним лежит следующий кадр. Без третьего индекса тот, кто сделает append(frame.Payload, ...), тихо перезапишет начало следующего кадра. С ним append выделит новый массив. В TS Buffer.subarray такой ловушки не имеет, потому что у буфера нет понятия ёмкости.

И честная пометка shortcut: в коде: длина не ограничена сверху. Клиент может прислать заголовок с длиной в 4 ГБ, и сервер начнёт копить байты, пока не кончится память. Промышленный сервер отвергает кадры больше разумного предела и закрывает соединение. Для учебного сервера на loopback это приемлемо, но знать об этом надо.

Сообщения

Теперь кодеки payload. Заведи wire/messages.go. Сначала типы.

// Байты типов сообщений (ASCII).
const (
	TypeQuery           byte = 'Q'
	TypeTerminate       byte = 'X'
	TypeRowDescription  byte = 'T'
	TypeDataRow         byte = 'D'
	TypeCommandComplete byte = 'C'
	TypeErrorResponse   byte = 'E'
	TypeReadyForQuery   byte = 'Z'
)

// Field это поле результата: имя колонки и её тип.
type Field = tuple.Column

var typeCodes = []tuple.ColumnType{tuple.TypeInt, tuple.TypeText, tuple.TypeBool}

// ClientMessage это Query или Terminate.
type ClientMessage interface{ clientType() byte }

type Query struct{ SQL string }
type Terminate struct{}

func (Query) clientType() byte     { return TypeQuery }
func (Terminate) clientType() byte { return TypeTerminate }

// ServerMessage это RowDescription, DataRow, CommandComplete, ErrorResponse или ReadyForQuery.
type ServerMessage interface{ serverType() byte }

type RowDescription struct{ Fields []Field }

// DataRow на стороне сервера несёт типизированные значения (int64, string, bool, nil), а
// после DecodeServer это сырой текст: string или nil. Тип отдаёт RowDescription, см.
// DecodeValueText.
type DataRow struct{ Values []tuple.Value }

type CommandComplete struct{ Tag string }

type ErrorResponse struct{ Code, Message string }

// ReadyForQuery несёт статус транзакции: 'I' idle, 'T' в транзакции, 'E' упавшая.
type ReadyForQuery struct{ Status byte }

func (RowDescription) serverType() byte  { return TypeRowDescription }
func (DataRow) serverType() byte         { return TypeDataRow }
func (CommandComplete) serverType() byte { return TypeCommandComplete }
func (ErrorResponse) serverType() byte   { return TypeErrorResponse }
func (ReadyForQuery) serverType() byte   { return TypeReadyForQuery }

Это Go-способ записать закрытую сумму типов. В TS было бы объединение { type: 'RowDescription', ... } | { type: 'DataRow', ... }. В Go объединений нет, зато есть интерфейс с неэкспортированным методом: serverType() со строчной буквы виден только внутри пакета wire, поэтому реализовать ServerMessage снаружи нельзя. Набор вариантов закрыт, и метод сразу несёт полезную нагрузку: байт типа кадра. Методы объявлены на получателе-значении без имени (func (Query) clientType()): самого получателя метод не читает.

type Field = tuple.Column со знаком = это псевдоним типа, а не новый тип: поле протокола и колонка схемы это одно и то же, и res.Fields из движка уходит в RowDescription без преобразования.

Тонкость в DataRow. Сервер строит его из типизированных значений, он знает типы. Клиент получает сырой текст и ещё не может его разобрать: типы он узнает только из пришедшего раньше RowDescription. Поэтому одна и та же структура на стороне сервера несёт int64, string, bool или nil, а после разбора только string или nil. Комментарий честно это фиксирует. Дальше кодирование:

// EncodeClient собирает кадр клиентского сообщения. Query это SQL плюс завершающий NUL.
func EncodeClient(m ClientMessage) []byte {
	if q, ok := m.(Query); ok {
		return EncodeFrame(TypeQuery, append([]byte(q.SQL), 0))
	}
	return EncodeFrame(m.clientType(), nil)
}

// EncodeServer собирает кадр серверного сообщения.
func EncodeServer(m ServerMessage) ([]byte, error) {
	var p []byte
	switch m := m.(type) {
	case RowDescription:
		p = be.AppendUint16(p, uint16(len(m.Fields)))
		for _, f := range m.Fields {
			code := -1
			for i, t := range typeCodes {
				if t == f.Type {
					code = i
				}
			}
			if code < 0 {
				return nil, fmt.Errorf("неизвестный тип колонки %q", f.Type)
			}
			p = append(be.AppendUint16(p, uint16(len(f.Name))), f.Name...)
			p = append(p, byte(code))
		}
	case DataRow:
		p = be.AppendUint16(p, uint16(len(m.Values)))
		for _, v := range m.Values {
			if v == nil {
				p = be.AppendUint32(p, 0xffffffff) // i32 -1 это NULL
				continue
			}
			text, err := EncodeValueText(v)
			if err != nil {
				return nil, err
			}
			p = append(be.AppendUint32(p, uint32(len(text))), text...)
		}
	case CommandComplete:
		p = []byte(m.Tag)
	case ErrorResponse:
		p = append(append(be.AppendUint16(p, uint16(len(m.Code))), m.Code...), m.Message...)
	case ReadyForQuery:
		p = []byte{m.Status}
	}
	return EncodeFrame(m.serverType(), p), nil
}

// EncodeValueText это текстовое представление значения в DataRow: int десятичной строкой,
// bool как t или f, text как есть.
func EncodeValueText(v tuple.Value) ([]byte, error) {
	switch v := v.(type) {
	case int64:
		return strconv.AppendInt(nil, v, 10), nil
	case bool:
		if v {
			return []byte("t"), nil
		}
		return []byte("f"), nil
	case string:
		return []byte(v), nil
	}
	return nil, fmt.Errorf("значение %T не кодируется в DataRow", v)
}

Query несёт SQL с завершающим нулевым байтом, как в Postgres; Terminate пуст. switch m := m.(type) затеняет m внутри каждой ветки переменной уже конкретного типа: в ветке RowDescription у m есть поле Fields. NULL в DataRow это длина 0xffffffff: на проводе u32, а по смыслу i32 минус один. strconv.AppendInt пишет число сразу в срез байт без промежуточной строки.

Разбор устроен зеркально, но payload приходит из сети, и каждой длине верить нельзя. Тут Go-порт применяет идиому из статьи Роба Пайка “Errors are values”:

// reader это курсор по payload с проверкой границ; первая нехватка байт залипает в err.
type reader struct {
	b   []byte
	err error
}

func (r *reader) bytes(n int) []byte {
	if r.err != nil || n < 0 || n > len(r.b) {
		r.err = errTruncated
		return nil
	}
	out := r.b[:n]
	r.b = r.b[n:]
	return out
}

func (r *reader) u16() uint16 {
	if b := r.bytes(2); b != nil {
		return be.Uint16(b)
	}
	return 0
}

Первая нехватка байт залипает в r.err, и все следующие чтения возвращают нули, ничего не трогая. Поэтому разбор DataRow пишется прямолинейно, без if err != nil после каждого поля:

	case TypeDataRow:
		n := int(r.u16())
		m := DataRow{Values: []tuple.Value{}}
		for i := 0; i < n && r.err == nil; i++ {
			length := int32(r.u32())
			if length == -1 {
				m.Values = append(m.Values, nil)
				continue
			}
			m.Values = append(m.Values, string(r.bytes(int(length))))
		}
		return m, r.err

Ошибка проверяется один раз, в конце, и в условии цикла, чтобы мусорное n = 65535 не крутило цикл впустую. int32(r.u32()) превращает 0xffffffff в минус один. А n < 0 в bytes ловит отрицательную длину, которая получилась бы из любого другого большого u32. Тест TestTruncatedPayloadIsError подсовывает DataRow, где длина ячейки 9, а байт один: разбор возвращает ошибку, а не падает с index out of range.

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

Кадры и сообщения готовы. Теперь логика: что сервер делает с Query. Сессия превращает SQL-строку в последовательность серверных сообщений и держит состояние транзакции этого соединения. Заведи wire/session.go.

// Session это состояние одного соединения: SQL-строка превращается в серверные сообщения,
// последнее из них всегда ReadyForQuery. Без сокетов, поэтому тестируется без сети.
//
// Каждое соединение это своя транзакция: BEGIN берёт собственный хэндл db.Begin(), а не
// общий sessionXid движка. Вне транзакции DML идёт в autocommit, SELECT берёт лёгкий снимок.
// Ошибка внутри BEGIN переводит транзакцию в failed: до ROLLBACK команды отвергаются.
type Session struct {
	db     *engine.Database
	tx     *engine.Transaction
	failed bool
}

// NewSession открывает сессию над базой.
func NewSession(db *engine.Database) *Session { return &Session{db: db} }

// Run выполняет SQL и возвращает ответ.
func (s *Session) Run(src string) []ServerMessage {
	cmd, err := sql.ParseCommand(src)
	if err != nil {
		return s.respond(errorOf(err))
	}
	if t, ok := cmd.(*sql.Txn); ok {
		return s.respond(s.txn(t.Op))
	}
	if s.failed {
		return s.respond(ErrorResponse{Code: "25P02", Message: "текущая транзакция прервана"})
	}
	var res engine.Result
	if s.tx == nil {
		res, err = s.db.Execute(src)
	} else {
		res, err = s.tx.Execute(src)
	}
	if err != nil {
		s.failed = s.tx != nil
		return s.respond(errorOf(err))
	}
	return s.respond(messagesFor(res)...)
}

// Close откатывает открытую транзакцию (Terminate или обрыв соединения).
func (s *Session) Close() error {
	if s.tx == nil {
		return nil
	}
	tx := s.tx
	s.tx, s.failed = nil, false
	return tx.Rollback()
}

func (s *Session) txn(op string) ServerMessage {
	switch {
	case op == "begin":
		if s.tx == nil {
			s.tx, s.failed = s.db.Begin(), false
		}
		return CommandComplete{Tag: "BEGIN"}
	case s.tx == nil:
		return CommandComplete{Tag: strings.ToUpper(op)}
	case op == "commit" && !s.failed:
		tx := s.tx
		s.tx = nil
		if err := tx.Commit(); err != nil {
			return errorOf(err)
		}
		return CommandComplete{Tag: "COMMIT"}
	}
	// ROLLBACK, а COMMIT упавшей транзакции это тоже откат
	if err := s.Close(); err != nil {
		return errorOf(err)
	}
	return CommandComplete{Tag: "ROLLBACK"}
}

func (s *Session) respond(msgs ...ServerMessage) []ServerMessage {
	status := byte('I')
	switch {
	case s.failed:
		status = 'E'
	case s.tx != nil:
		status = 'T'
	}
	return append(msgs, ReadyForQuery{Status: status})
}

Ключевое решение: каждое соединение это своя транзакция. BEGIN берёт собственный хэндл db.Begin() из урока 8, а не сессионную транзакцию движка. Поле db.session у Database одно на весь процесс, и если бы сетевой BEGIN шёл через db.Execute("BEGIN"), все соединения делили бы одну транзакцию. Вне транзакции сессия зовёт db.Execute: DML идёт в autocommit, SELECT берёт лёгкий снимок.

Обработка ошибок повторяет Postgres. Оператор упал внутри транзакции, и сессия переходит в failed (статус E). До ROLLBACK любая команда отвергается кодом 25P02, а COMMIT упавшей транзакции превращается в откат и честно отвечает тегом ROLLBACK. BEGIN и COMMIT вне транзакции не ошибка, а пустая операция с тегом: так ведёт себя и psql. Close откатывает незакрытую транзакцию: это сработает при Terminate и при обрыве соединения и не даст утечь локам строк.

Сессия не знает про сокеты. Run берёт строку и отдаёт срез сообщений, поэтому вся логика протокола тестируется без сети, обычными юнит-тестами. Это та же граница, что в TS, и в Go она даётся даром: сервер ниже это тонкий слой ввода-вывода вокруг Run.

Две свободные функции завершают сессию:

func messagesFor(res engine.Result) []ServerMessage {
	switch res.Kind {
	case engine.ResultSelect:
		out := []ServerMessage{RowDescription{Fields: res.Fields}}
		for _, row := range res.Rows {
			out = append(out, DataRow{Values: row})
		}
		return append(out, CommandComplete{Tag: fmt.Sprintf("SELECT %d", len(res.Rows))})
	case engine.ResultExplain:
		return []ServerMessage{
			RowDescription{Fields: []Field{{Name: "QUERY PLAN", Type: tuple.TypeText}}},
			DataRow{Values: []tuple.Value{res.Plan}},
			CommandComplete{Tag: "EXPLAIN"},
		}
	}
	return []ServerMessage{CommandComplete{Tag: res.Tag}}
}

// errorOf даёт SQLSTATE-подобный код по классу ошибки: синтаксис 42601, ошибки движка со
// своим кодом, всё прочее (ввод-вывод, поломка) XX000.
func errorOf(err error) ServerMessage {
	var syntax *sql.Error
	code := "XX000"
	if errors.As(err, &syntax) {
		code = "42601"
	} else if c := dberr.CodeOf(err); c != "" {
		code = string(c)
	}
	return ErrorResponse{Code: code, Message: err.Error()}
}

messagesFor переводит результат движка на язык протокола. EXPLAIN притворяется однострочным SELECT с колонкой QUERY PLAN, как в Postgres. А errorOf показывает, ради чего в уроке 8 ошибки получили коды: TS-сессии нужна цепочка instanceof в правильном порядке (подклассы раньше базового), а здесь это одна строка dberr.CodeOf. Синтаксическая ошибка это отдельный тип *sql.Error со смещением, его ищет errors.As. Всё прочее, например ошибка ввода-вывода, это XX000: клиенту незачем знать подробности.