rubydb 0.1.4 → 0.1.6

data/accelerator/internal/parallel/worker_pool.go

@@ -0,0 +1,53 @@
+package parallel
+
+import (
+	"context"
+	"sync"
+)
+
+type WorkerPool struct {
+	queue  *BoundedQueue[func()]
+	wait   sync.WaitGroup
+	closed chan struct{}
+}
+
+func NewWorkerPool(workers, queueSize int) *WorkerPool {
+	if workers < 1 {
+		workers = 1
+	}
+	pool := &WorkerPool{queue: NewBoundedQueue[func()](queueSize), closed: make(chan struct{})}
+	for index := 0; index < workers; index++ {
+		pool.wait.Add(1)
+		go func() {
+			defer pool.wait.Done()
+			for {
+				job, err := pool.queue.Pop(context.Background())
+				if err != nil {
+					return
+				}
+				job()
+			}
+		}()
+	}
+	return pool
+}
+
+func (pool *WorkerPool) Submit(ctx context.Context, job func()) error {
+	select {
+	case <-pool.closed:
+		return ErrQueueClosed
+	default:
+	}
+	return pool.queue.Push(ctx, job)
+}
+
+func (pool *WorkerPool) Close() {
+	select {
+	case <-pool.closed:
+		return
+	default:
+		close(pool.closed)
+		pool.queue.Close()
+		pool.wait.Wait()
+	}
+}

data/accelerator/internal/protocol/cancellation.go

@@ -0,0 +1,48 @@
+package protocol
+
+import (
+	"context"
+	"sync"
+)
+
+// CancellationRegistry gives the runtime a bounded, request-scoped context
+// registry. A future multiplexed transport can cancel a running request by
+// ID without changing the physical execution APIs.
+type CancellationRegistry struct {
+	mu      sync.Mutex
+	entries map[string]context.CancelFunc
+}
+
+func NewCancellationRegistry() *CancellationRegistry {
+	return &CancellationRegistry{entries: make(map[string]context.CancelFunc)}
+}
+
+func (registry *CancellationRegistry) Begin(parent context.Context, id string) context.Context {
+	ctx, cancel := context.WithCancel(parent)
+	registry.mu.Lock()
+	if previous, exists := registry.entries[id]; exists {
+		previous()
+	}
+	registry.entries[id] = cancel
+	registry.mu.Unlock()
+	return ctx
+}
+
+func (registry *CancellationRegistry) Cancel(id string) bool {
+	registry.mu.Lock()
+	cancel, exists := registry.entries[id]
+	if exists {
+		delete(registry.entries, id)
+	}
+	registry.mu.Unlock()
+	if exists {
+		cancel()
+	}
+	return exists
+}
+
+func (registry *CancellationRegistry) Finish(id string) {
+	registry.mu.Lock()
+	delete(registry.entries, id)
+	registry.mu.Unlock()
+}

data/accelerator/internal/protocol/columnar.go

@@ -0,0 +1,263 @@
+package protocol
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"errors"
+	"math"
+)
+
+func EncodeRows(rows []map[string]interface{}) ([]byte, error) {
+	columns := make([]string, 0)
+	known := make(map[string]bool)
+	for _, row := range rows {
+		for key := range row {
+			if !known[key] {
+				known[key] = true
+				columns = append(columns, key)
+			}
+		}
+	}
+	if len(columns) > math.MaxUint16 || uint64(len(rows)) > math.MaxUint32 {
+		return nil, errors.New("binary rows batch is too large")
+	}
+	output := bytes.NewBuffer(make([]byte, 0, 6+len(columns)*16))
+	writeUint16(output, uint16(len(columns)))
+	writeUint32(output, uint32(len(rows)))
+	for _, column := range columns {
+		if len(column) > math.MaxUint16 {
+			return nil, errors.New("binary row column name is too long")
+		}
+		writeUint16(output, uint16(len(column)))
+		output.WriteString(column)
+	}
+	for _, column := range columns {
+		for _, row := range rows {
+			if err := encodeValue(output, row[column]); err != nil {
+				return nil, err
+			}
+		}
+	}
+	return output.Bytes(), nil
+}
+
+func DecodeRows(data []byte) ([]map[string]interface{}, error) {
+	reader := newValueReader(data)
+	columnCount, err := reader.uint16()
+	if err != nil {
+		return nil, errors.New("binary rows batch is truncated")
+	}
+	rowCount, err := reader.uint32()
+	if err != nil {
+		return nil, errors.New("binary rows batch is truncated")
+	}
+	columns := make([]string, columnCount)
+	for index := range columns {
+		name, err := reader.string16()
+		if err != nil {
+			return nil, errors.New("binary row column metadata is invalid")
+		}
+		columns[index] = name
+	}
+	rows := make([]map[string]interface{}, rowCount)
+	for index := range rows {
+		rows[index] = make(map[string]interface{}, columnCount)
+	}
+	for _, column := range columns {
+		for rowIndex := range rows {
+			value, err := reader.value()
+			if err != nil {
+				return nil, errors.New("binary row value is invalid")
+			}
+			rows[rowIndex][column] = value
+		}
+	}
+	if reader.remaining() != 0 {
+		return nil, errors.New("binary rows batch has trailing bytes")
+	}
+	return rows, nil
+}
+
+func encodeValue(output *bytes.Buffer, value interface{}) error {
+	switch typed := value.(type) {
+	case nil:
+		output.WriteByte(0)
+	case bool:
+		output.WriteByte(1)
+		if typed {
+			output.WriteByte(1)
+		} else {
+			output.WriteByte(0)
+		}
+	case json.Number:
+		if integer, err := typed.Int64(); err == nil {
+			output.WriteByte(2)
+			writeInt64(output, integer)
+		} else if number, err := typed.Float64(); err == nil {
+			output.WriteByte(3)
+			writeFloat64(output, number)
+		} else {
+			return errors.New("invalid numeric value")
+		}
+	case int:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case int8:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case int16:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case int32:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case int64:
+		output.WriteByte(2)
+		writeInt64(output, typed)
+	case uint:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case uint8:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case uint16:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case uint32:
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case uint64:
+		if typed > math.MaxInt64 {
+			return errors.New("unsigned numeric value is too large")
+		}
+		output.WriteByte(2)
+		writeInt64(output, int64(typed))
+	case float32:
+		output.WriteByte(3)
+		writeFloat64(output, float64(typed))
+	case float64:
+		output.WriteByte(3)
+		writeFloat64(output, typed)
+	case []byte:
+		output.WriteByte(5)
+		writeBytes(output, typed)
+	case string:
+		output.WriteByte(4)
+		writeBytes(output, []byte(typed))
+	default:
+		encoded, err := json.Marshal(value)
+		if err != nil {
+			return errors.New("value is not serializable")
+		}
+		output.WriteByte(6)
+		writeBytes(output, encoded)
+	}
+	return nil
+}
+
+type valueReader struct {
+	data   []byte
+	offset int
+}
+
+func newValueReader(data []byte) *valueReader { return &valueReader{data: data} }
+func (reader *valueReader) remaining() int    { return len(reader.data) - reader.offset }
+
+func (reader *valueReader) take(length int) ([]byte, error) {
+	if length < 0 || reader.remaining() < length {
+		return nil, errors.New("binary value is truncated")
+	}
+	value := reader.data[reader.offset : reader.offset+length]
+	reader.offset += length
+	return value, nil
+}
+
+func (reader *valueReader) uint16() (uint16, error) {
+	value, err := reader.take(2)
+	if err != nil {
+		return 0, err
+	}
+	return binary.LittleEndian.Uint16(value), nil
+}
+
+func (reader *valueReader) uint32() (uint32, error) {
+	value, err := reader.take(4)
+	if err != nil {
+		return 0, err
+	}
+	return binary.LittleEndian.Uint32(value), nil
+}
+
+func (reader *valueReader) string16() (string, error) {
+	length, err := reader.uint16()
+	if err != nil {
+		return "", err
+	}
+	value, err := reader.take(int(length))
+	return string(value), err
+}
+
+func (reader *valueReader) value() (interface{}, error) {
+	tag, err := reader.take(1)
+	if err != nil {
+		return nil, err
+	}
+	switch tag[0] {
+	case 0:
+		return nil, nil
+	case 1:
+		value, err := reader.take(1)
+		return value[0] == 1, err
+	case 2:
+		value, err := reader.take(8)
+		if err != nil {
+			return nil, err
+		}
+		return int64(binary.LittleEndian.Uint64(value)), nil
+	case 3:
+		value, err := reader.take(8)
+		if err != nil {
+			return nil, err
+		}
+		return math.Float64frombits(binary.LittleEndian.Uint64(value)), nil
+	case 4, 5, 6:
+		length, err := reader.uint32()
+		if err != nil {
+			return nil, err
+		}
+		value, err := reader.take(int(length))
+		if err != nil {
+			return nil, err
+		}
+		if tag[0] == 4 {
+			return string(value), nil
+		}
+		if tag[0] == 5 {
+			return append([]byte(nil), value...), nil
+		}
+		var decoded interface{}
+		if err := json.Unmarshal(value, &decoded); err != nil {
+			return nil, err
+		}
+		return decoded, nil
+	default:
+		return nil, errors.New("unknown binary value tag")
+	}
+}
+
+func writeUint16(output *bytes.Buffer, value uint16) {
+	_ = binary.Write(output, binary.LittleEndian, value)
+}
+func writeUint32(output *bytes.Buffer, value uint32) {
+	_ = binary.Write(output, binary.LittleEndian, value)
+}
+func writeBytes(output *bytes.Buffer, value []byte) {
+	writeUint32(output, uint32(len(value)))
+	output.Write(value)
+}
+func writeInt64(output *bytes.Buffer, value int64)     { writeUint64(output, uint64(value)) }
+func writeFloat64(output *bytes.Buffer, value float64) { writeUint64(output, math.Float64bits(value)) }
+func writeUint64(output *bytes.Buffer, value uint64) {
+	_ = binary.Write(output, binary.LittleEndian, value)
+}

data/accelerator/internal/protocol/frame.go

@@ -0,0 +1,187 @@
+package protocol
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+)
+
+const (
+	Version       = 1
+	BinaryVersion = 1
+	MaxFrameSize  = 16 * 1024 * 1024
+
+	BinaryRequestKind      byte = 1
+	BinaryResponseKind     byte = 2
+	BinaryRowsType         byte = 1
+	BinaryJoinType         byte = 2
+	BinarySnapshotScanType byte = 3
+)
+
+var BinaryMagic = []byte("RDBB")
+
+var ErrFrameTooLarge = errors.New("frame too large")
+
+type BinaryFrame struct {
+	Kind    byte
+	TypeID  byte
+	Status  byte
+	ID      string
+	Payload []byte
+}
+
+func ReadMessage(reader *bufio.Reader) (*BinaryFrame, []byte, error) {
+	prefix, err := reader.Peek(len(BinaryMagic))
+	if err != nil {
+		return nil, nil, err
+	}
+	if bytes.Equal(prefix, BinaryMagic) {
+		frame, err := ReadBinaryFrame(reader)
+		return frame, nil, err
+	}
+	frame, err := ReadFrame(reader)
+	return nil, frame, err
+}
+
+func ReadBinaryFrame(reader *bufio.Reader) (*BinaryFrame, error) {
+	header := make([]byte, 14)
+	if _, err := io.ReadFull(reader, header); err != nil {
+		return nil, err
+	}
+	if !bytes.Equal(header[:4], BinaryMagic) || header[4] != BinaryVersion {
+		return nil, errors.New("invalid binary frame")
+	}
+	idLength := int(binaryLittleEndian16(header[8:10]))
+	payloadLength := int(binaryLittleEndian32(header[10:14]))
+	if payloadLength > MaxFrameSize || idLength > 1024 || payloadLength+idLength > MaxFrameSize {
+		return nil, ErrFrameTooLarge
+	}
+	body := make([]byte, idLength+payloadLength)
+	if _, err := io.ReadFull(reader, body); err != nil {
+		return nil, err
+	}
+	return &BinaryFrame{
+		Kind:    header[5],
+		TypeID:  header[6],
+		Status:  header[7],
+		ID:      string(body[:idLength]),
+		Payload: body[idLength:],
+	}, nil
+}
+
+func WriteBinaryResponse(writer *bufio.Writer, requestFrame BinaryFrame, payload []byte, message, code string) error {
+	status := byte(0)
+	if message != "" {
+		status = 1
+		payload = MustJSON(map[string]string{"error": message, "code": code})
+	}
+	if len(requestFrame.ID) > 1024 || len(payload)+len(requestFrame.ID) > MaxFrameSize {
+		return ErrFrameTooLarge
+	}
+	header := make([]byte, 14)
+	copy(header[:4], BinaryMagic)
+	header[4] = BinaryVersion
+	header[5] = BinaryResponseKind
+	header[6] = requestFrame.TypeID
+	header[7] = status
+	putBinaryLittleEndian16(header[8:10], uint16(len(requestFrame.ID)))
+	putBinaryLittleEndian32(header[10:14], uint32(len(payload)))
+	if _, err := writer.Write(header); err != nil {
+		return err
+	}
+	if _, err := writer.WriteString(requestFrame.ID); err != nil {
+		return err
+	}
+	if _, err := writer.Write(payload); err != nil {
+		return err
+	}
+	return writer.Flush()
+}
+
+func ReadFrame(reader *bufio.Reader) ([]byte, error) {
+	frame := make([]byte, 0, 1024)
+	for {
+		chunk, err := reader.ReadSlice('\n')
+		if len(frame)+len(chunk) > MaxFrameSize {
+			return nil, ErrFrameTooLarge
+		}
+		frame = append(frame, chunk...)
+		if err == bufio.ErrBufferFull {
+			continue
+		}
+		if err != nil {
+			return nil, err
+		}
+		return frame, nil
+	}
+}
+
+func WriteJSONResponse(writer *bufio.Writer, value interface{}) error {
+	data, err := json.Marshal(value)
+	if err != nil {
+		return err
+	}
+	if len(data)+1 > MaxFrameSize {
+		return fmt.Errorf("response exceeds maximum frame size")
+	}
+	if _, err = writer.Write(append(data, '\n')); err != nil {
+		return err
+	}
+	return writer.Flush()
+}
+
+func MustJSON(value interface{}) []byte {
+	data, _ := json.Marshal(value)
+	return data
+}
+
+type Reader struct {
+	data   []byte
+	offset int
+}
+
+func NewReader(data []byte) *Reader { return &Reader{data: data} }
+
+func (reader *Reader) Remaining() int { return len(reader.data) - reader.offset }
+
+func (reader *Reader) Bytes() []byte { return reader.data[reader.offset:] }
+
+func (reader *Reader) Take(length int) ([]byte, error) {
+	if length < 0 || reader.Remaining() < length {
+		return nil, errors.New("binary value is truncated")
+	}
+	value := reader.data[reader.offset : reader.offset+length]
+	reader.offset += length
+	return value, nil
+}
+
+func (reader *Reader) Uint32() (uint32, error) {
+	value, err := reader.Take(4)
+	if err != nil {
+		return 0, err
+	}
+	return binaryLittleEndian32(value), nil
+}
+
+func binaryLittleEndian16(data []byte) uint16 {
+	return uint16(data[0]) | uint16(data[1])<<8
+}
+
+func binaryLittleEndian32(data []byte) uint32 {
+	return uint32(data[0]) | uint32(data[1])<<8 | uint32(data[2])<<16 | uint32(data[3])<<24
+}
+
+func putBinaryLittleEndian16(data []byte, value uint16) {
+	data[0] = byte(value)
+	data[1] = byte(value >> 8)
+}
+
+func putBinaryLittleEndian32(data []byte, value uint32) {
+	data[0] = byte(value)
+	data[1] = byte(value >> 8)
+	data[2] = byte(value >> 16)
+	data[3] = byte(value >> 24)
+}