pizzasql/pkg/storage/kv.go

package storage

import (
	"bufio"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

const (
	headerSize    = 32
	headerMagic   = "PKBF"
	headerVersion = 1

	opPing         = 1
	opStatus       = 2
	opGet          = 3
	opPut          = 4
	opDelete       = 5
	opExists       = 6
	opMultiGet     = 7
	opBatchWrite   = 8
	opScanOpen     = 9
	opScanNext     = 10
	opScanClose    = 11
	opCompareBatch = 12

	batchPut    = 1
	batchDelete = 2

	statusOK       = 0
	statusNotFound = 1
	statusError    = 2

	maxKeySize         = 1024 * 1024
	maxValueSize       = 64 * 1024 * 1024
	maxTransactionSize = 64 * 1024 * 1024
	maxOperations      = 65535
	maxFrameSize       = maxKeySize + maxValueSize + 1024

	scanPageSize       = 1024
	existsPipelineSize = 128
)

var crc32cTable = crc32.MakeTable(crc32.Castagnoli)

var (
	ErrKeyNotFound = errors.New("key not found")
	ErrProtocol    = errors.New("pkbfi protocol error")
)

func crc32c(p []byte) uint32 {
	return crc32.Checksum(p, crc32cTable)
}

func putU16(b []byte, v uint16) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
}

func putU32(b []byte, v uint32) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
	b[2] = byte(v >> 16)
	b[3] = byte(v >> 24)
}

func putU64(b []byte, v uint64) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
	b[2] = byte(v >> 16)
	b[3] = byte(v >> 24)
	b[4] = byte(v >> 32)
	b[5] = byte(v >> 40)
	b[6] = byte(v >> 48)
	b[7] = byte(v >> 56)
}

func getU16(b []byte) uint16 {
	return uint16(b[0]) | uint16(b[1])<<8
}

func getU32(b []byte) uint32 {
	return uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16 | uint32(b[3])<<24
}

func getU64(b []byte) uint64 {
	return uint64(b[0]) | uint64(b[1])<<8 | uint64(b[2])<<16 | uint64(b[3])<<24 |
		uint64(b[4])<<32 | uint64(b[5])<<40 | uint64(b[6])<<48 | uint64(b[7])<<56
}

func encodeFrame(opcode, flags uint16, requestID uint64, payload []byte) []byte {
	frame := make([]byte, headerSize+len(payload))
	copy(frame[0:4], headerMagic)
	putU16(frame[4:6], headerVersion)
	putU16(frame[6:8], 0)
	putU16(frame[8:10], opcode)
	putU16(frame[10:12], flags)
	putU64(frame[12:20], requestID)
	putU32(frame[20:24], uint32(len(payload)))
	putU32(frame[24:28], crc32c(payload))
	putU32(frame[28:32], 0)
	putU32(frame[28:32], crc32c(frame[0:32]))
	copy(frame[32:], payload)
	return frame
}

func readFrame(r *bufio.Reader) (uint16, uint16, uint64, []byte, error) {
	var header [headerSize]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return 0, 0, 0, nil, err
	}
	if string(header[0:4]) != headerMagic {
		return 0, 0, 0, nil, fmt.Errorf("%w: invalid magic", ErrProtocol)
	}
	if getU16(header[4:6]) != headerVersion {
		return 0, 0, 0, nil, fmt.Errorf("%w: incompatible version", ErrProtocol)
	}
	payloadLen := getU32(header[20:24])
	if payloadLen > maxFrameSize {
		return 0, 0, 0, nil, fmt.Errorf("%w: frame too large", ErrProtocol)
	}
	headerCRC := getU32(header[28:32])
	var headerCopy [headerSize]byte
	copy(headerCopy[:], header[:])
	putU32(headerCopy[28:32], 0)
	if crc32c(headerCopy[:]) != headerCRC {
		return 0, 0, 0, nil, fmt.Errorf("%w: header checksum mismatch", ErrProtocol)
	}
	payload := make([]byte, payloadLen)
	if _, err := io.ReadFull(r, payload); err != nil {
		return 0, 0, 0, nil, err
	}
	if crc32c(payload) != getU32(header[24:28]) {
		return 0, 0, 0, nil, fmt.Errorf("%w: payload checksum mismatch", ErrProtocol)
	}
	return getU16(header[8:10]), getU16(header[10:12]), getU64(header[12:20]), payload, nil
}

func encodeResponse(opcode uint16, requestID uint64, body []byte) []byte {
	return encodeFrame(opcode|0x8000, 1, requestID, body)
}

func oneKeyPayload(key []byte) []byte {
	payload := make([]byte, 4+len(key))
	putU32(payload[0:4], uint32(len(key)))
	copy(payload[4:], key)
	return payload
}

func validateKey(key []byte) error {
	if len(key) > maxKeySize {
		return fmt.Errorf("pkbfi: key exceeds %d bytes", maxKeySize)
	}
	return nil
}

func validateValue(value []byte) error {
	if len(value) > maxValueSize {
		return fmt.Errorf("pkbfi: value exceeds %d bytes", maxValueSize)
	}
	return nil
}

func parseOneKey(payload []byte) ([]byte, bool) {
	if len(payload) < 4 {
		return nil, false
	}
	length := getU32(payload[0:4])
	if length > maxKeySize || uint64(4)+uint64(length) != uint64(len(payload)) {
		return nil, false
	}
	return payload[4:], true
}

func errorBody(message string) []byte {
	body := make([]byte, 2+len(message))
	putU16(body[0:2], statusError)
	copy(body[2:], message)
	return body
}

type KVClient struct {
	conn           net.Conn
	reader         *bufio.Reader
	writer         *bufio.Writer
	mu             sync.Mutex
	nextID         uint64
	requestTimeout time.Duration
	lastUsed       time.Time
}

type KVResult struct {
	Value []byte
	LSN   uint64
	Found bool
}

type KVEntry struct {
	Key   []byte
	Value []byte
	LSN   uint64
}

type BatchOp struct {
	Op    byte
	Key   []byte
	Value []byte
}

type CompareCheck struct {
	Key []byte
	LSN uint64
}

type ScanCursor struct {
	client *KVClient
	id     uint64
}

func NewKVClient(addr string) (*KVClient, error) {
	network, target := parseAddr(addr)
	conn, err := net.Dial(network, target)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to PizzaKV: %w", err)
	}
	return &KVClient{
		conn:     conn,
		reader:   bufio.NewReader(conn),
		writer:   bufio.NewWriter(conn),
		nextID:   1,
		lastUsed: time.Now(),
	}, nil
}

func parseAddr(addr string) (string, string) {
	if strings.HasPrefix(addr, "unix:") {
		return "unix", strings.TrimPrefix(addr, "unix:")
	}
	return "tcp", addr
}

func (c *KVClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn != nil {
		err := c.conn.Close()
		c.conn = nil
		return err
	}
	return nil
}

func (c *KVClient) SetDeadline(t time.Time) error {
	if c.conn == nil {
		return nil
	}
	return c.conn.SetDeadline(t)
}

func (c *KVClient) writeFrame(opcode, flags uint16, requestID uint64, payload []byte) error {
	frame := encodeFrame(opcode, flags, requestID, payload)
	if _, err := c.writer.Write(frame); err != nil {
		return err
	}
	return c.writer.Flush()
}

func (c *KVClient) request(opcode uint16, payload []byte) (uint16, []byte, error) {
	if c.requestTimeout > 0 {
		if err := c.conn.SetDeadline(time.Now().Add(c.requestTimeout)); err != nil {
			return 0, nil, err
		}
	}
	requestID := c.nextID
	c.nextID++
	if err := c.writeFrame(opcode, 0, requestID, payload); err != nil {
		return 0, nil, err
	}
	respOpcode, respFlags, respID, body, err := readFrame(c.reader)
	if err != nil {
		return 0, nil, err
	}
	if respOpcode != opcode|0x8000 {
		return 0, nil, fmt.Errorf("%w: unexpected response opcode %d", ErrProtocol, respOpcode)
	}
	if respFlags != 1 {
		return 0, nil, fmt.Errorf("%w: unexpected response flags %d", ErrProtocol, respFlags)
	}
	if respID != requestID {
		return 0, nil, fmt.Errorf("%w: response id %d does not match request %d", ErrProtocol, respID, requestID)
	}
	c.lastUsed = time.Now()
	if len(body) < 2 {
		return 0, nil, fmt.Errorf("%w: response too short", ErrProtocol)
	}
	status := getU16(body[0:2])
	if status == statusError {
		return status, nil, fmt.Errorf("pkbfi server error: %s", body[2:])
	}
	return status, body[2:], nil
}

func (c *KVClient) Put(key, value []byte) (uint64, error) {
	if err := validateKey(key); err != nil {
		return 0, err
	}
	if err := validateValue(value); err != nil {
		return 0, err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload := make([]byte, 8+len(key)+len(value))
	putU32(payload[0:4], uint32(len(key)))
	putU32(payload[4:8], uint32(len(value)))
	copy(payload[8:], key)
	copy(payload[8+len(key):], value)
	status, body, err := c.request(opPut, payload)
	if err != nil {
		return 0, err
	}
	if status != statusOK || len(body) != 8 {
		return 0, fmt.Errorf("%w: malformed put response", ErrProtocol)
	}
	return getU64(body[0:8]), nil
}

func (c *KVClient) Get(key []byte) (KVResult, error) {
	if err := validateKey(key); err != nil {
		return KVResult{}, err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	status, body, err := c.request(opGet, oneKeyPayload(key))
	if err != nil {
		return KVResult{}, err
	}
	if status == statusNotFound {
		return KVResult{}, ErrKeyNotFound
	}
	if status != statusOK || len(body) < 12 {
		return KVResult{}, fmt.Errorf("%w: malformed get response", ErrProtocol)
	}
	lsn := getU64(body[0:8])
	valueLen := getU32(body[8:12])
	if valueLen > maxValueSize || uint64(len(body)) != 12+uint64(valueLen) {
		return KVResult{}, fmt.Errorf("%w: malformed get value length", ErrProtocol)
	}
	return KVResult{Value: body[12:], LSN: lsn, Found: true}, nil
}

func (c *KVClient) Del(key []byte) (bool, error) {
	if err := validateKey(key); err != nil {
		return false, err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	status, body, err := c.request(opDelete, oneKeyPayload(key))
	if err != nil {
		return false, err
	}
	if status != statusOK || len(body) != 1 {
		return false, fmt.Errorf("%w: malformed delete response", ErrProtocol)
	}
	return body[0] != 0, nil
}

func (c *KVClient) Exists(key []byte) (bool, error) {
	if err := validateKey(key); err != nil {
		return false, err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	status, body, err := c.request(opExists, oneKeyPayload(key))
	if err != nil {
		return false, err
	}
	if status != statusOK || len(body) != 1 {
		return false, fmt.Errorf("%w: malformed exists response", ErrProtocol)
	}
	return body[0] != 0, nil
}

func (c *KVClient) ExistsMany(keys [][]byte) ([]bool, error) {
	if len(keys) > maxOperations {
		return nil, fmt.Errorf("pkbfi: too many keys")
	}
	for _, key := range keys {
		if err := validateKey(key); err != nil {
			return nil, err
		}
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	results := make([]bool, len(keys))
	for start := 0; start < len(keys); start += existsPipelineSize {
		end := start + existsPipelineSize
		if end > len(keys) {
			end = len(keys)
		}
		ids := make([]uint64, end-start)
		if c.requestTimeout > 0 {
			if err := c.conn.SetDeadline(time.Now().Add(c.requestTimeout)); err != nil {
				return nil, err
			}
		}
		for i, key := range keys[start:end] {
			ids[i] = c.nextID
			c.nextID++
			if _, err := c.writer.Write(encodeFrame(opExists, 0, ids[i], oneKeyPayload(key))); err != nil {
				return nil, err
			}
		}
		if err := c.writer.Flush(); err != nil {
			return nil, err
		}
		for i, requestID := range ids {
			opcode, flags, responseID, body, err := readFrame(c.reader)
			if err != nil {
				return nil, err
			}
			if opcode != opExists|0x8000 || flags != 1 || responseID != requestID {
				return nil, fmt.Errorf("%w: malformed exists response frame", ErrProtocol)
			}
			if len(body) < 2 {
				return nil, fmt.Errorf("%w: response too short", ErrProtocol)
			}
			status := getU16(body[0:2])
			if status == statusError {
				return nil, fmt.Errorf("pkbfi server error: %s", body[2:])
			}
			if status != statusOK || len(body) != 3 {
				return nil, fmt.Errorf("%w: malformed exists response", ErrProtocol)
			}
			results[start+i] = body[2] != 0
			c.lastUsed = time.Now()
		}
	}
	return results, nil
}

func (c *KVClient) MultiGet(keys [][]byte) ([]KVResult, error) {
	if len(keys) > maxOperations {
		return nil, fmt.Errorf("pkbfi: too many keys")
	}
	payloadSize := 4
	for _, key := range keys {
		if err := validateKey(key); err != nil {
			return nil, err
		}
		payloadSize += 4 + len(key)
		if payloadSize > maxFrameSize {
			return nil, fmt.Errorf("pkbfi: multi_get request exceeds frame limit")
		}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload := make([]byte, 4, payloadSize)
	putU32(payload[0:4], uint32(len(keys)))
	for _, key := range keys {
		var length [4]byte
		putU32(length[:], uint32(len(key)))
		payload = append(payload, length[:]...)
		payload = append(payload, key...)
	}
	status, body, err := c.request(opMultiGet, payload)
	if err != nil {
		return nil, err
	}
	if status != statusOK || len(body) < 4 {
		return nil, fmt.Errorf("%w: malformed multi_get response", ErrProtocol)
	}
	count := getU32(body[0:4])
	if count != uint32(len(keys)) {
		return nil, fmt.Errorf("%w: multi_get count mismatch", ErrProtocol)
	}
	results := make([]KVResult, count)
	pos := 4
	for i := uint32(0); i < count; i++ {
		if len(body)-pos < 16 {
			return nil, fmt.Errorf("%w: truncated multi_get entry", ErrProtocol)
		}
		present := body[pos] != 0
		valueLen := getU32(body[pos+4 : pos+8])
		lsn := getU64(body[pos+8 : pos+16])
		pos += 16
		results[i] = KVResult{LSN: lsn, Found: present}
		if present {
			if valueLen > maxValueSize || len(body)-pos < int(valueLen) {
				return nil, fmt.Errorf("%w: multi_get value length", ErrProtocol)
			}
			results[i].Value = body[pos : pos+int(valueLen)]
			pos += int(valueLen)
		}
	}
	if pos != len(body) {
		return nil, fmt.Errorf("%w: multi_get trailing bytes", ErrProtocol)
	}
	return results, nil
}

func (c *KVClient) BatchWrite(ops []BatchOp, metadata []byte) (uint64, error) {
	if len(ops) == 0 || len(ops) > maxOperations {
		return 0, fmt.Errorf("pkbfi: invalid operation count")
	}
	if len(metadata) > maxTransactionSize-8 {
		return 0, fmt.Errorf("pkbfi: batch metadata exceeds transaction limit")
	}
	payloadSize := 8 + len(metadata)
	for _, op := range ops {
		if op.Op != batchPut && op.Op != batchDelete {
			return 0, fmt.Errorf("pkbfi: invalid batch opcode %d", op.Op)
		}
		if err := validateKey(op.Key); err != nil {
			return 0, err
		}
		if err := validateValue(op.Value); err != nil {
			return 0, err
		}
		if op.Op == batchDelete && len(op.Value) != 0 {
			return 0, fmt.Errorf("pkbfi: delete operation with value")
		}
		payloadSize += 12 + len(op.Key) + len(op.Value)
		if payloadSize > maxTransactionSize {
			return 0, fmt.Errorf("pkbfi: batch exceeds transaction limit")
		}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload := make([]byte, 8, payloadSize)
	putU32(payload[0:4], uint32(len(ops)))
	putU32(payload[4:8], uint32(len(metadata)))
	payload = append(payload, metadata...)
	for _, op := range ops {
		var header [12]byte
		header[0] = op.Op
		putU32(header[4:8], uint32(len(op.Key)))
		putU32(header[8:12], uint32(len(op.Value)))
		payload = append(payload, header[:]...)
		payload = append(payload, op.Key...)
		payload = append(payload, op.Value...)
	}
	status, body, err := c.request(opBatchWrite, payload)
	if err != nil {
		return 0, err
	}
	if status != statusOK || len(body) != 8 {
		return 0, fmt.Errorf("%w: malformed batch_write response", ErrProtocol)
	}
	return getU64(body[0:8]), nil
}

func (c *KVClient) CompareBatchWrite(checks []CompareCheck, ops []BatchOp, metadata []byte) (uint64, bool, error) {
	if len(ops) == 0 || len(ops) > maxOperations {
		return 0, false, fmt.Errorf("pkbfi: invalid operation count")
	}
	if len(checks) > maxOperations {
		return 0, false, fmt.Errorf("pkbfi: too many compare checks")
	}
	if len(metadata) > maxTransactionSize-16 {
		return 0, false, fmt.Errorf("pkbfi: batch metadata exceeds transaction limit")
	}
	payloadSize := 16 + len(metadata)
	for _, check := range checks {
		if err := validateKey(check.Key); err != nil {
			return 0, false, err
		}
		payloadSize += 16 + len(check.Key)
		if payloadSize > maxFrameSize {
			return 0, false, fmt.Errorf("pkbfi: compare batch exceeds frame limit")
		}
	}
	for _, op := range ops {
		if op.Op != batchPut && op.Op != batchDelete {
			return 0, false, fmt.Errorf("pkbfi: invalid batch opcode %d", op.Op)
		}
		if err := validateKey(op.Key); err != nil {
			return 0, false, err
		}
		if err := validateValue(op.Value); err != nil {
			return 0, false, err
		}
		if op.Op == batchDelete && len(op.Value) != 0 {
			return 0, false, fmt.Errorf("pkbfi: delete operation with value")
		}
		payloadSize += 12 + len(op.Key) + len(op.Value)
		if payloadSize > maxTransactionSize {
			return 0, false, fmt.Errorf("pkbfi: batch exceeds transaction limit")
		}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload := make([]byte, 16, payloadSize)
	putU32(payload[0:4], uint32(len(checks)))
	putU32(payload[4:8], uint32(len(ops)))
	putU32(payload[8:12], uint32(len(metadata)))
	for _, check := range checks {
		var header [16]byte
		putU32(header[0:4], uint32(len(check.Key)))
		putU64(header[8:16], check.LSN)
		payload = append(payload, header[:]...)
		payload = append(payload, check.Key...)
	}
	payload = append(payload, metadata...)
	for _, op := range ops {
		var header [12]byte
		header[0] = op.Op
		putU32(header[4:8], uint32(len(op.Key)))
		putU32(header[8:12], uint32(len(op.Value)))
		payload = append(payload, header[:]...)
		payload = append(payload, op.Key...)
		payload = append(payload, op.Value...)
	}
	status, body, err := c.request(opCompareBatch, payload)
	if err != nil {
		return 0, false, err
	}
	if status != statusOK || len(body) != 16 {
		return 0, false, fmt.Errorf("%w: malformed compare_batch_write response", ErrProtocol)
	}
	committed := body[0] != 0
	lsn := getU64(body[8:16])
	return lsn, committed, nil
}

func (c *KVClient) Scan(prefix []byte) (*ScanCursor, error) {
	return c.openScan(prefix, true, scanPageSize)
}

func (c *KVClient) ScanWithLimit(prefix []byte, pageSize uint32) (*ScanCursor, error) {
	return c.openScan(prefix, true, pageSize)
}

// ScanKeys opens a key-only scan: the server omits values from the returned
// pages. It is used when only the key set (or its size) is needed, such as the
// COUNT(*) fast path or bulk deletion, so a full-table scan does not pull row
// values across the wire.
func (c *KVClient) ScanKeys(prefix []byte) (*ScanCursor, error) {
	return c.openScan(prefix, false, scanPageSize)
}

func (c *KVClient) openScan(prefix []byte, includeValues bool, pageSize uint32) (*ScanCursor, error) {
	if err := validateKey(prefix); err != nil {
		return nil, err
	}
	if pageSize == 0 || pageSize > 4096 {
		return nil, fmt.Errorf("pkbfi: scan page size must be between 1 and 4096")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload := make([]byte, 12+len(prefix))
	if includeValues {
		payload[0] = 1
	}
	putU32(payload[4:8], pageSize)
	putU32(payload[8:12], uint32(len(prefix)))
	copy(payload[12:], prefix)
	status, body, err := c.request(opScanOpen, payload)
	if err != nil {
		return nil, err
	}
	if status != statusOK || len(body) != 8 {
		return nil, fmt.Errorf("%w: malformed scan_open response", ErrProtocol)
	}
	return &ScanCursor{client: c, id: getU64(body[0:8])}, nil
}

func (s *ScanCursor) Next() ([]KVEntry, bool, error) {
	s.client.mu.Lock()
	defer s.client.mu.Unlock()
	var payload [12]byte
	putU64(payload[0:8], s.id)
	putU32(payload[8:12], 0)
	status, body, err := s.client.request(opScanNext, payload[:])
	if err != nil {
		return nil, false, err
	}
	if status != statusOK || len(body) < 8 {
		return nil, false, fmt.Errorf("%w: malformed scan_next response", ErrProtocol)
	}
	done := body[0] != 0
	count := getU32(body[4:8])
	entries := make([]KVEntry, 0, count)
	pos := 8
	for i := uint32(0); i < count; i++ {
		if len(body)-pos < 16 {
			return nil, false, fmt.Errorf("%w: truncated scan entry", ErrProtocol)
		}
		keyLen := getU32(body[pos : pos+4])
		valueLen := getU32(body[pos+4 : pos+8])
		lsn := getU64(body[pos+8 : pos+16])
		pos += 16
		if keyLen > maxKeySize || valueLen > maxValueSize || len(body)-pos < int(keyLen)+int(valueLen) {
			return nil, false, fmt.Errorf("%w: scan entry length", ErrProtocol)
		}
		key := body[pos : pos+int(keyLen)]
		pos += int(keyLen)
		value := body[pos : pos+int(valueLen)]
		pos += int(valueLen)
		entries = append(entries, KVEntry{Key: key, Value: value, LSN: lsn})
	}
	if pos != len(body) {
		return nil, false, fmt.Errorf("%w: scan trailing bytes", ErrProtocol)
	}
	return entries, done, nil
}

func (s *ScanCursor) Close() error {
	s.client.mu.Lock()
	defer s.client.mu.Unlock()
	var payload [8]byte
	putU64(payload[0:8], s.id)
	status, body, err := s.client.request(opScanClose, payload[:])
	if err != nil {
		return err
	}
	if status != statusOK || len(body) != 1 {
		return fmt.Errorf("%w: malformed scan_close response", ErrProtocol)
	}
	return nil
}

func (c *KVClient) Write(key, value string) error {
	_, err := c.Put([]byte(key), []byte(value))
	return err
}

func (c *KVClient) Read(key string) (string, error) {
	res, err := c.Get([]byte(key))
	if err != nil {
		return "", err
	}
	return string(res.Value), nil
}

func (c *KVClient) Delete(key string) error {
	_, err := c.Del([]byte(key))
	return err
}

func (c *KVClient) Reads(prefix string) ([]string, error) {
	scan, err := c.Scan([]byte(prefix))
	if err != nil {
		return nil, err
	}
	defer scan.Close()
	values := make([]string, 0)
	for {
		entries, done, err := scan.Next()
		if err != nil {
			return nil, err
		}
		for _, entry := range entries {
			values = append(values, string(entry.Value))
		}
		if done {
			return values, nil
		}
	}
}

func (c *KVClient) IsAlive() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn == nil {
		return false
	}
	c.conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
	defer c.conn.SetDeadline(time.Time{})
	_, _, err := c.request(opPing, nil)
	return err == nil
}

type KVPool struct {
	addr    string
	pool    chan *KVClient
	size    int
	timeout time.Duration
	mu      sync.Mutex
	closed  bool
}

func (p *KVPool) replacementClient() (*KVClient, error) {
	client, err := NewKVClient(p.addr)
	if err == nil {
		client.requestTimeout = p.timeout
		return client, nil
	}
	p.mu.Lock()
	if !p.closed {
		select {
		case p.pool <- nil:
		default:
		}
	}
	p.mu.Unlock()
	return nil, err
}

func NewKVPool(addr string, size int, timeout time.Duration) (*KVPool, error) {
	p := &KVPool{
		addr:    addr,
		pool:    make(chan *KVClient, size),
		size:    size,
		timeout: timeout,
	}
	for i := 0; i < size; i++ {
		client, err := NewKVClient(addr)
		if err != nil {
			p.Close()
			return nil, fmt.Errorf("failed to create connection pool: %w", err)
		}
		p.pool <- client
	}
	return p, nil
}

func (p *KVPool) Get() (*KVClient, error) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil, fmt.Errorf("pool is closed")
	}
	p.mu.Unlock()

	select {
	case client := <-p.pool:
		if client != nil && client.conn != nil {
			if client.lastUsed.IsZero() {
				client.lastUsed = time.Now()
			} else if time.Since(client.lastUsed) >= 20*time.Second && !client.IsAlive() {
				client.Close()
				return p.replacementClient()
			}
			client.requestTimeout = p.timeout
			return client, nil
		}
		return p.replacementClient()
	case <-time.After(30 * time.Second):
		return nil, fmt.Errorf("kv pool timeout: no connection available after 30s")
	}
}

func (p *KVPool) Put(client *KVClient) {
	if client == nil {
		return
	}

	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		client.Close()
		return
	}
	p.mu.Unlock()

	client.requestTimeout = 0
	client.SetDeadline(time.Time{})

	select {
	case p.pool <- client:
	default:
		client.Close()
	}
}

func (p *KVPool) Close() error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	p.closed = true
	p.mu.Unlock()

	close(p.pool)
	for client := range p.pool {
		if client != nil {
			client.Close()
		}
	}
	return nil
}

func (p *KVPool) WithClient(fn func(*KVClient) error) error {
	client, err := p.Get()
	if err != nil {
		return err
	}
	if err := fn(client); err != nil {
		if !isConnectionError(err) {
			p.Put(client)
			return err
		}
		client.Close()
		p.mu.Lock()
		if !p.closed {
			select {
			case p.pool <- nil:
			default:
			}
		}
		p.mu.Unlock()
		return err
	}
	p.Put(client)
	return nil
}

func isConnectionError(err error) bool {
	var netErr net.Error
	if errors.As(err, &netErr) {
		return true
	}
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
		return true
	}
	return errors.Is(err, ErrProtocol)
}