From d410df2d7e8d2ece9295c76558caae09ac76948b Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Sat, 8 Jul 2023 19:08:28 +0200
Subject: [PATCH] Improved data safety

---
 Collection_test.go |  2 +-
 FileStorage.go     | 17 ++++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/Collection_test.go b/Collection_test.go
index c160b33..394224a 100644
--- a/Collection_test.go
+++ b/Collection_test.go
@@ -198,7 +198,7 @@ func BenchmarkColdStart(b *testing.B) {
 	defer users.Sync()
 	defer users.Clear()
 
-	b.Run("100 records", func(b *testing.B) {
+	b.Run("100", func(b *testing.B) {
 		for i := 0; i < 100; i++ {
 			users.Set(strconv.Itoa(i), &User{Name: fmt.Sprintf("User %d", i)})
 		}
diff --git a/FileStorage.go b/FileStorage.go
index ca57c69..7c59f27 100644
--- a/FileStorage.go
+++ b/FileStorage.go
@@ -7,9 +7,9 @@ import (
 	"log"
 	"os"
 	"path/filepath"
-	"runtime"
 	"sort"
 	"sync/atomic"
+	"time"
 )
 
 type FileStorage[T any] struct {
@@ -55,7 +55,7 @@ func (fs *FileStorage[T]) Sync() {
 
 func (fs *FileStorage[T]) flushWorker() {
 	for {
-		runtime.Gosched()
+		time.Sleep(time.Millisecond)
 
 		if fs.dirty.Swap(0) == 0 {
 			select {
@@ -75,8 +75,9 @@ func (fs *FileStorage[T]) flushWorker() {
 }
 
 func (fs *FileStorage[T]) flush() error {
-	fileName := filepath.Join(fs.collection.root, fs.collection.name+".dat")
-	file, err := os.OpenFile(fileName, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	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
@@ -104,7 +105,13 @@ func (fs *FileStorage[T]) flush() error {
 		return err
 	}
 
-	return file.Close()
+	err = file.Close()
+
+	if err != nil {
+		return err
+	}
+
+	return os.Rename(newPath, oldPath)
 }
 
 // readFrom reads the entire collection.