go-kv/wal/block_writer.go

package wal

import (
	"errors"
	"fmt"
	"io"
)

// BlockWriter manages a single 32 KB block buffer for writing physical records.
// It handles block boundary padding and flushing complete blocks to an io.Writer.
type BlockWriter struct {
	buf    [WalBlockSize]byte
	offset uint32 // current write position within the block
}

// NewBlockWriter creates a BlockWriter ready to write into a fresh block.
func NewBlockWriter() *BlockWriter {
	return &BlockWriter{}
}

// BlockOffset returns the current write offset within the block (0..WalBlockSize).
func (bw *BlockWriter) BlockOffset() uint32 {
	return bw.offset
}

// WriteRecord writes a single physical record into the block buffer.
// If the record (header + payload) does not fit in the remaining space,
// the current block is padded with zeros and flushed to w, then the record
// is written at the start of a fresh block.
//
// Precondition: payload length must be ≤ WalBlockSize - PhysicalRecordHeaderSize
// (the caller is responsible for splitting large batches into appropriately-sized chunks).
func (bw *BlockWriter) WriteRecord(recType uint8, payload []byte, w io.Writer) error {
	recordSize := PhysicalRecordHeaderSize + len(payload)

	if recordSize > WalBlockSize {
		return fmt.Errorf("wal: record size %d exceeds block size %d",
			recordSize, WalBlockSize)
	}

	// Check if padding is needed before writing this record.
	pad := bw.paddingNeeded()
	if pad > 0 {
		// Pad remaining bytes with zeros and flush.
		if err := bw.flushPadded(w, pad); err != nil {
			return fmt.Errorf("wal: flushing padded block: %w", err)
		}
	}

	// Check if the record fits in the current block.
	remaining := WalBlockSize - bw.offset
	if uint32(recordSize) > remaining {
		// Not enough room — pad the rest and flush, then start a new block.
		pad = int(remaining)
		if err := bw.flushPadded(w, pad); err != nil {
			return fmt.Errorf("wal: flushing partial block: %w", err)
		}
	}

	// Encode physical record directly into the block buffer.
	encoded := EncodePhysicalRecord(recType, payload)
	copy(bw.buf[bw.offset:], encoded)
	bw.offset += uint32(len(encoded))

	// If the block is exactly full, flush it immediately.
	if bw.offset == WalBlockSize {
		if _, err := w.Write(bw.buf[:]); err != nil {
			return fmt.Errorf("wal: writing full block: %w", err)
		}
		bw.offset = 0
	}

	return nil
}

// Flush writes the current block buffer to w, padding unused bytes with zeros.
// If the block is empty (offset == 0), this is a no-op.
func (bw *BlockWriter) Flush(w io.Writer) error {
	if bw.offset == 0 {
		return nil
	}
	return bw.flushPadded(w, int(WalBlockSize-bw.offset))
}

// Reset clears the block buffer, returning it to an empty state.
func (bw *BlockWriter) Reset() {
	bw.offset = 0
	// Zero the buffer so partial blocks are padded with zeros.
	for i := range bw.buf {
		bw.buf[i] = 0
	}
}

// paddingNeeded returns the number of zero-padding bytes required at the current
// block offset. When the remaining space in the block is ≤ PhysicalRecordHeaderSize (7),
// that space cannot hold even a minimal physical record and must be zero-padded.
func (bw *BlockWriter) paddingNeeded() int {
	remaining := WalBlockSize - bw.offset
	if remaining <= PhysicalRecordHeaderSize {
		return int(remaining)
	}
	return 0
}

// flushPadded pads the remaining bytes with zeros and writes the full block to w.
// pad is the number of trailing bytes to zero-fill (WalBlockSize - offset - pad already zero
// from initial state or previous Reset).
func (bw *BlockWriter) flushPadded(w io.Writer, pad int) error {
	if pad <= 0 {
		return nil
	}

	// Zero-fill padding region. The buffer was zeroed at init/reset,
	// but we write explicitly for safety after partial record writes.
	for i := uint32(0); i < uint32(pad); i++ {
		bw.buf[bw.offset+i] = 0
	}

	if _, err := w.Write(bw.buf[:]); err != nil {
		return err
	}

	bw.offset = 0
	for i := range bw.buf {
		bw.buf[i] = 0
	}

	return nil
}

// Bytes returns a copy of the current block contents up to the current offset.
// Useful for testing.
func (bw *BlockWriter) Bytes() []byte {
	out := make([]byte, bw.offset)
	copy(out, bw.buf[:bw.offset])
	return out
}

// FullBlockBytes returns the full block buffer. Only valid when offset == WalBlockSize.
func (bw *BlockWriter) FullBlockBytes() []byte {
	out := make([]byte, WalBlockSize)
	copy(out, bw.buf[:])
	return out
}

// errBlockWriterNil is returned when a nil writer is passed to write operations.
var errBlockWriterNil = errors.New("wal: writer must not be nil")