193 lines
2.9 KiB
Go
193 lines
2.9 KiB
Go
package ocean
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/json"
|
|
"io"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type FileStorage[T any] struct {
|
|
collection *collection[T]
|
|
dirty atomic.Uint32
|
|
sync chan struct{}
|
|
}
|
|
|
|
func (fs *FileStorage[T]) Init(c *collection[T]) error {
|
|
fs.collection = c
|
|
fs.sync = make(chan struct{})
|
|
|
|
go fs.flushWorker()
|
|
|
|
fileName := filepath.Join(c.root, c.name+".dat")
|
|
file, err := os.OpenFile(fileName, os.O_RDONLY, 0600)
|
|
|
|
if os.IsNotExist(err) {
|
|
return nil
|
|
}
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
defer file.Close()
|
|
return fs.readFrom(file)
|
|
}
|
|
|
|
func (fs *FileStorage[T]) Delete(key string) error {
|
|
fs.dirty.Store(1)
|
|
return nil
|
|
}
|
|
|
|
func (fs *FileStorage[T]) Set(key string, value *T) error {
|
|
fs.dirty.Store(1)
|
|
return nil
|
|
}
|
|
|
|
func (fs *FileStorage[T]) Sync() {
|
|
<-fs.sync
|
|
}
|
|
|
|
func (fs *FileStorage[T]) flushWorker() {
|
|
for {
|
|
time.Sleep(time.Millisecond)
|
|
|
|
if fs.dirty.Swap(0) == 0 {
|
|
select {
|
|
case fs.sync <- struct{}{}:
|
|
default:
|
|
}
|
|
|
|
continue
|
|
}
|
|
|
|
err := fs.flush()
|
|
|
|
if err != nil {
|
|
log.Println(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (fs *FileStorage[T]) flush() error {
|
|
oldPath := filepath.Join(fs.collection.root, fs.collection.name+".dat")
|
|
newPath := oldPath + ".tmp"
|
|
file, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
bufferedWriter := bufio.NewWriter(file)
|
|
err = fs.writeTo(bufferedWriter)
|
|
|
|
if err != nil {
|
|
file.Close()
|
|
return err
|
|
}
|
|
|
|
err = bufferedWriter.Flush()
|
|
|
|
if err != nil {
|
|
file.Close()
|
|
return err
|
|
}
|
|
|
|
err = file.Sync()
|
|
|
|
if err != nil {
|
|
file.Close()
|
|
return err
|
|
}
|
|
|
|
err = file.Close()
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return os.Rename(newPath, oldPath)
|
|
}
|
|
|
|
// readFrom reads the entire collection.
|
|
func (fs *FileStorage[T]) readFrom(stream io.Reader) error {
|
|
var (
|
|
key string
|
|
value []byte
|
|
)
|
|
|
|
scanner := bufio.NewScanner(stream)
|
|
|
|
for scanner.Scan() {
|
|
if key == "" {
|
|
key = scanner.Text()
|
|
continue
|
|
}
|
|
|
|
value = scanner.Bytes()
|
|
object := new(T)
|
|
err := json.Unmarshal(value, object)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
fs.collection.data.Store(key, object)
|
|
key = ""
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// writeTo writes the entire collection.
|
|
func (fs *FileStorage[T]) writeTo(writer io.Writer) error {
|
|
stringWriter, ok := writer.(io.StringWriter)
|
|
|
|
if !ok {
|
|
panic("The given io.Writer is not an io.StringWriter")
|
|
}
|
|
|
|
records := []keyValue{}
|
|
|
|
fs.collection.data.Range(func(key, value any) bool {
|
|
records = append(records, keyValue{
|
|
key: key.(string),
|
|
value: value,
|
|
})
|
|
return true
|
|
})
|
|
|
|
sort.Slice(records, func(i, j int) bool {
|
|
return records[i].key < records[j].key
|
|
})
|
|
|
|
encoder := NewEncoder(writer)
|
|
|
|
for _, record := range records {
|
|
_, err := stringWriter.WriteString(record.key)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = stringWriter.WriteString("\n")
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = encoder.Encode(record.value)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|