go-kv/wal/segment_writer.go

package wal

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/dailz/go-kv/config"
)

// SegmentWriter handles appending WAL batches to a single segment file.
// It manages block-aligned writes via BlockWriter and tracks file offset
// for segment rotation decisions.
type SegmentWriter struct {
	fd            *os.File
	dir           string
	cfg           *config.WalConfig
	segmentID     uint64
	startSequence uint64
	blockWriter   *BlockWriter
	currentOffset uint64 // total bytes written (starts at WalFileHeaderSize)
	maxPayload    uint64 // cfg.MaxSegmentSize - WalFileHeaderSize
}

// NewSegmentWriter creates a new WAL segment file and writes the file header.
// The segment file is created with a .tmp extension, the header is written and
// synced, then the file is atomically renamed to its final name and synced again.
func NewSegmentWriter(
	dir string,
	segmentID uint64,
	startSequence uint64,
	cfg *config.WalConfig,
) (*SegmentWriter, error) {
	if err := cfg.Validate(); err != nil {
		return nil, fmt.Errorf("wal: invalid config: %w", err)
	}

	baseName := fmt.Sprintf("segment-%d.wal", segmentID)
	tmpPath := filepath.Join(dir, baseName+".tmp")
	finalPath := filepath.Join(dir, baseName)

	// Create the temp file.
	fd, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o644)
	if err != nil {
		return nil, fmt.Errorf("wal: create segment temp file %s: %w", tmpPath, err)
	}

	// Build and write the file header.
	hdr := &WalFileHeader{
		BlockSize:     cfg.BlockSize,
		SegmentID:     segmentID,
		StartSequence: startSequence,
	}
	encoded := EncodeWalHeader(hdr)

	if _, err := fd.Write(encoded[:]); err != nil {
		fd.Close()
		os.Remove(tmpPath)
		return nil, fmt.Errorf("wal: write segment header: %w", err)
	}

	// Sync the header to disk.
	if err := fd.Sync(); err != nil {
		fd.Close()
		os.Remove(tmpPath)
		return nil, fmt.Errorf("wal: sync segment header: %w", err)
	}

	// Atomically rename temp file to final name.
	if err := fd.Close(); err != nil {
		os.Remove(tmpPath)
		return nil, fmt.Errorf("wal: close temp file: %w", err)
	}

	if err := os.Rename(tmpPath, finalPath); err != nil {
		os.Remove(tmpPath)
		return nil, fmt.Errorf("wal: rename segment file: %w", err)
	}

	// Open the final file for appending.
	fd, err = os.OpenFile(finalPath, os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return nil, fmt.Errorf("wal: open segment file for append: %w", err)
	}

	// Sync directory to make rename durable (best-effort on Linux).
	if dirFD, derr := os.Open(dir); derr == nil {
		dirFD.Sync()
		dirFD.Close()
	}

	maxPayload := cfg.MaxSegmentSize - WalFileHeaderSize

	return &SegmentWriter{
		fd:            fd,
		dir:           dir,
		cfg:           cfg,
		segmentID:     segmentID,
		startSequence: startSequence,
		blockWriter:   NewBlockWriter(),
		currentOffset: WalFileHeaderSize,
		maxPayload:    maxPayload,
	}, nil
}

// AppendBatch encodes the batch into physical records and appends them to the
// segment file. The encoded batch is split into block-aligned physical records
// using SplitIntoRecords.
func (sw *SegmentWriter) AppendBatch(encodedBatch []byte) error {
	records := SplitIntoRecords(encodedBatch)
	if len(records) == 0 {
		return nil
	}

	for _, rec := range records {
		if len(rec) < PhysicalRecordHeaderSize {
			return fmt.Errorf("wal: corrupted physical record: size %d < header size %d",
				len(rec), PhysicalRecordHeaderSize)
		}

		recType := rec[6] // type byte is at offset 6 in the encoded record
		payload := rec[PhysicalRecordHeaderSize:]

		if err := sw.blockWriter.WriteRecord(recType, payload, sw.fd); err != nil {
			return fmt.Errorf("wal: writing physical record: %w", err)
		}

		sw.currentOffset += uint64(len(rec))
	}

	return nil
}

// Sync flushes the segment file to durable storage.
func (sw *SegmentWriter) Sync() error {
	return sw.fd.Sync()
}

// Close flushes any partial block and closes the segment file.
func (sw *SegmentWriter) Close() error {
	if err := sw.blockWriter.Flush(sw.fd); err != nil {
		return fmt.Errorf("wal: flushing block writer on close: %w", err)
	}
	return sw.fd.Close()
}

// RemainingPayload returns the number of bytes that can still be written
// to this segment before it reaches its maximum size.
func (sw *SegmentWriter) RemainingPayload() uint64 {
	if sw.currentOffset >= sw.cfg.MaxSegmentSize {
		return 0
	}
	return sw.cfg.MaxSegmentSize - sw.currentOffset
}

// CurrentOffset returns the total number of bytes written to the segment file,
// including the file header.
func (sw *SegmentWriter) CurrentOffset() uint64 {
	return sw.currentOffset
}

// SegmentID returns the segment identifier.
func (sw *SegmentWriter) SegmentID() uint64 {
	return sw.segmentID
}

// SegmentPath returns the full filesystem path to the segment file.
func (sw *SegmentWriter) SegmentPath() string {
	return filepath.Join(sw.dir, fmt.Sprintf("segment-%d.wal", sw.segmentID))
}