package sstable import ( "bytes" "encoding/binary" "fmt" "os" "path/filepath" "github.com/jeremytregunna/kevo/pkg/sstable/block" "github.com/jeremytregunna/kevo/pkg/sstable/footer" ) // FileManager handles file operations for SSTable writing type FileManager struct { path string tmpPath string file *os.File } // NewFileManager creates a new FileManager for the given file path func NewFileManager(path string) (*FileManager, error) { // Create temporary file for writing dir := filepath.Dir(path) tmpPath := filepath.Join(dir, fmt.Sprintf(".%s.tmp", filepath.Base(path))) file, err := os.Create(tmpPath) if err != nil { return nil, fmt.Errorf("failed to create temporary file: %w", err) } return &FileManager{ path: path, tmpPath: tmpPath, file: file, }, nil } // Write writes data to the file at the current position func (fm *FileManager) Write(data []byte) (int, error) { return fm.file.Write(data) } // Sync flushes the file to disk func (fm *FileManager) Sync() error { return fm.file.Sync() } // Close closes the file func (fm *FileManager) Close() error { if fm.file == nil { return nil } err := fm.file.Close() fm.file = nil return err } // FinalizeFile closes the file and renames it to the final path func (fm *FileManager) FinalizeFile() error { // Close the file before renaming if err := fm.Close(); err != nil { return fmt.Errorf("failed to close file: %w", err) } // Rename the temp file to the final path if err := os.Rename(fm.tmpPath, fm.path); err != nil { return fmt.Errorf("failed to rename temp file: %w", err) } return nil } // Cleanup removes the temporary file if writing is aborted func (fm *FileManager) Cleanup() error { if fm.file != nil { fm.Close() } return os.Remove(fm.tmpPath) } // BlockManager handles block building and serialization type BlockManager struct { builder *block.Builder offset uint64 } // NewBlockManager creates a new BlockManager func NewBlockManager() *BlockManager { return &BlockManager{ builder: block.NewBuilder(), offset: 0, } } // Add adds a key-value pair to the current block func (bm *BlockManager) Add(key, value []byte) error { return bm.builder.Add(key, value) } // EstimatedSize returns the estimated size of the current block func (bm *BlockManager) EstimatedSize() uint32 { return bm.builder.EstimatedSize() } // Entries returns the number of entries in the current block func (bm *BlockManager) Entries() int { return bm.builder.Entries() } // GetEntries returns all entries in the current block func (bm *BlockManager) GetEntries() []block.Entry { return bm.builder.GetEntries() } // Reset resets the block builder func (bm *BlockManager) Reset() { bm.builder.Reset() } // Serialize serializes the current block func (bm *BlockManager) Serialize() ([]byte, error) { var buf bytes.Buffer _, err := bm.builder.Finish(&buf) if err != nil { return nil, fmt.Errorf("failed to finish block: %w", err) } return buf.Bytes(), nil } // IndexBuilder constructs the index block type IndexBuilder struct { builder *block.Builder entries []*IndexEntry } // NewIndexBuilder creates a new IndexBuilder func NewIndexBuilder() *IndexBuilder { return &IndexBuilder{ builder: block.NewBuilder(), entries: make([]*IndexEntry, 0), } } // AddIndexEntry adds an entry to the pending index entries func (ib *IndexBuilder) AddIndexEntry(entry *IndexEntry) { ib.entries = append(ib.entries, entry) } // BuildIndex builds the index block from the collected entries func (ib *IndexBuilder) BuildIndex() error { // Add all index entries to the index block for _, entry := range ib.entries { // Index entry format: key=firstKey, value=blockOffset+blockSize var valueBuf bytes.Buffer binary.Write(&valueBuf, binary.LittleEndian, entry.BlockOffset) binary.Write(&valueBuf, binary.LittleEndian, entry.BlockSize) if err := ib.builder.Add(entry.FirstKey, valueBuf.Bytes()); err != nil { return fmt.Errorf("failed to add index entry: %w", err) } } return nil } // Serialize serializes the index block func (ib *IndexBuilder) Serialize() ([]byte, error) { var buf bytes.Buffer _, err := ib.builder.Finish(&buf) if err != nil { return nil, fmt.Errorf("failed to finish index block: %w", err) } return buf.Bytes(), nil } // Writer writes an SSTable file type Writer struct { fileManager *FileManager blockManager *BlockManager indexBuilder *IndexBuilder dataOffset uint64 firstKey []byte lastKey []byte entriesAdded uint32 } // NewWriter creates a new SSTable writer func NewWriter(path string) (*Writer, error) { fileManager, err := NewFileManager(path) if err != nil { return nil, err } return &Writer{ fileManager: fileManager, blockManager: NewBlockManager(), indexBuilder: NewIndexBuilder(), dataOffset: 0, entriesAdded: 0, }, nil } // Add adds a key-value pair to the SSTable // Keys must be added in sorted order func (w *Writer) Add(key, value []byte) error { // Keep track of first and last keys if w.entriesAdded == 0 { w.firstKey = append([]byte(nil), key...) } w.lastKey = append([]byte(nil), key...) // Add to block if err := w.blockManager.Add(key, value); err != nil { return fmt.Errorf("failed to add to block: %w", err) } w.entriesAdded++ // Flush the block if it's getting too large // Use IndexKeyInterval to determine when to flush based on accumulated data size if w.blockManager.EstimatedSize() >= IndexKeyInterval { if err := w.flushBlock(); err != nil { return err } } return nil } // AddTombstone adds a deletion marker (tombstone) for a key to the SSTable // This is functionally equivalent to Add(key, nil) but makes the intention explicit func (w *Writer) AddTombstone(key []byte) error { return w.Add(key, nil) } // flushBlock writes the current block to the file and adds an index entry func (w *Writer) flushBlock() error { // Skip if the block is empty if w.blockManager.Entries() == 0 { return nil } // Record the offset of this block blockOffset := w.dataOffset // Get first key entries := w.blockManager.GetEntries() if len(entries) == 0 { return fmt.Errorf("block has no entries") } firstKey := entries[0].Key // Serialize the block blockData, err := w.blockManager.Serialize() if err != nil { return err } blockSize := uint32(len(blockData)) // Write the block to file n, err := w.fileManager.Write(blockData) if err != nil { return fmt.Errorf("failed to write block to file: %w", err) } if n != len(blockData) { return fmt.Errorf("wrote incomplete block: %d of %d bytes", n, len(blockData)) } // Add the index entry w.indexBuilder.AddIndexEntry(&IndexEntry{ BlockOffset: blockOffset, BlockSize: blockSize, FirstKey: firstKey, }) // Update offset for next block w.dataOffset += uint64(n) // Reset the block builder for next block w.blockManager.Reset() return nil } // Finish completes the SSTable writing process func (w *Writer) Finish() error { defer func() { w.fileManager.Close() }() // Flush any pending data block (only if we have entries that haven't been flushed) if w.blockManager.Entries() > 0 { if err := w.flushBlock(); err != nil { return err } } // Create index block indexOffset := w.dataOffset // Build the index from collected entries if err := w.indexBuilder.BuildIndex(); err != nil { return err } // Serialize and write the index block indexData, err := w.indexBuilder.Serialize() if err != nil { return err } indexSize := uint32(len(indexData)) n, err := w.fileManager.Write(indexData) if err != nil { return fmt.Errorf("failed to write index block: %w", err) } if n != len(indexData) { return fmt.Errorf("wrote incomplete index block: %d of %d bytes", n, len(indexData)) } // Create footer ft := footer.NewFooter( indexOffset, indexSize, w.entriesAdded, 0, // MinKeyOffset - not implemented yet 0, // MaxKeyOffset - not implemented yet ) // Serialize footer footerData := ft.Encode() // Write footer n, err = w.fileManager.Write(footerData) if err != nil { return fmt.Errorf("failed to write footer: %w", err) } if n != len(footerData) { return fmt.Errorf("wrote incomplete footer: %d of %d bytes", n, len(footerData)) } // Sync the file if err := w.fileManager.Sync(); err != nil { return fmt.Errorf("failed to sync file: %w", err) } // Finalize file (close and rename) return w.fileManager.FinalizeFile() } // Abort cancels the SSTable writing process func (w *Writer) Abort() error { return w.fileManager.Cleanup() }