tinydb/db.go

210 lines
4.6 KiB
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
}