tinydb/db.go

package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"sync"
)

var ErrClosed = errors.New("database has been closed")

type DB struct {
	sync.RWMutex                   // locking
	data         map[string]string // current state of the database
	file         *os.File          // data storage
}

// OpenDB opens an existing database file, loading the contents to memory.
func OpenDB(path string) (*DB, error) {
	file, err := os.OpenFile(path, os.O_RDWR, 0 /* not used */)
	if err != nil {
		return nil, err
	}

	// TODO: We might want a recover flag that just reads as much as it can
	// instead of returning io.ErrUnexpectedEOF.
	db := &DB{data: make(map[string]string), file: file}
	for {
		var header struct{ KeyLen, ValueLen int32 }
		err := binary.Read(db.file, binary.LittleEndian, &header)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		if header.KeyLen < 0 {
			return nil, fmt.Errorf("invalid key length: %d", header.KeyLen)
		}
		key := make([]byte, header.KeyLen)
		if n, err := file.Read(key); err != nil {
			return nil, err
		} else if n != len(key) {
			return nil, io.ErrUnexpectedEOF
		}

		if header.ValueLen < 0 {
			delete(db.data, string(key))
			continue
		}
		value := make([]byte, header.ValueLen)
		if n, err := file.Read(value); err != nil {
			return nil, err
		} else if n != len(value) {
			return nil, io.ErrUnexpectedEOF
		}

		db.data[string(key)] = string(value)
	}
	// We've been successful, clean up after failed snapshots
	os.Remove(path + ".1")
	return db, nil
}

// CreateDB creates a new database, overwriting any previous file contents.
func CreateDB(path string) (*DB, error) {
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return nil, err
	}
	return &DB{data: make(map[string]string), file: file}, nil
}

// Get retrieves the value corresponding to the given key.
func (db *DB) Get(key string) (string, bool, error) {
	db.RLock()
	defer db.RUnlock()

	if db.file == nil {
		return "", false, ErrClosed
	}

	value, ok := db.data[key]
	return value, ok, nil
}

func put(file *os.File, key, value string) error {
	header := [2]int32{int32(len(key)), int32(len(value))}
	if err := binary.Write(file, binary.LittleEndian, &header); err != nil {
		return err
	}

	if n, err := file.WriteString(key); err != nil {
		return err
	} else if n != len(key) {
		return io.ErrShortWrite
	}

	if n, err := file.WriteString(value); err != nil {
		return err
	} else if n != len(value) {
		return io.ErrShortWrite
	}
	return nil
}

// Put saves a key-value pair in the database storage.
func (db *DB) Put(key, value string) error {
	db.Lock()
	defer db.Unlock()

	if db.file == nil {
		return ErrClosed
	}

	// XXX we should check whether the key and the value fit
	if err := put(db.file, key, value); err != nil {
		return err
	}
	if err := db.file.Sync(); err != nil {
		return err
	}

	db.data[key] = value
	return nil
}

// Delete deletes a key from the database storage.
func (db *DB) Delete(key string) error {
	db.Lock()
	defer db.Unlock()

	if db.file == nil {
		return ErrClosed
	}

	// Like put(), just without the "value"
	header := [2]int32{int32(len(key)), -1}
	if err := binary.Write(db.file, binary.LittleEndian, &header); err != nil {
		return err
	}

	if n, err := db.file.WriteString(key); err != nil {
		return err
	} else if n != len(key) {
		return io.ErrShortWrite
	}

	if err := db.file.Sync(); err != nil {
		return err
	}

	// TODO maybe return an indication whether anything was actually deleted
	delete(db.data, key)
	return nil
}

// Checkpoint gets rid of historical data in the database file.
func (db *DB) Checkpoint() error {
	db.Lock()
	defer db.Unlock()

	if db.file == nil {
		return ErrClosed
	}

	checkpoint, err := os.OpenFile(db.file.Name()+".1",
		os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	// The checkpoint is made of the current value of every present key
	for key, value := range db.data {
		if err := put(checkpoint, key, value); err != nil {
			return err
		}
	}
	if err := checkpoint.Sync(); err != nil {
		return err
	}

	// Atomically move the checkpoint over the old database file
	if err := os.Rename(checkpoint.Name(), db.file.Name()); err != nil {
		// Not sure how much sense this makes--when do we get here?
		_ = os.Remove(checkpoint.Name())
		return err
	}
	// The old file now points to unlinked storage, replace it with the new one
	_ = db.file.Close()
	db.file = checkpoint
	return nil
}

// Close closes the database file, rendering the object unusable.
func (db *DB) Close() error {
	db.Lock()
	defer db.Unlock()

	if db.file != nil {
		return nil
	}

	err := db.file.Close()
	db.file = nil
	return err
}