add in support for batched writes

flatfs.go

@@ -68,12 +68,8 @@ func (fs *Datastore) decode(file string) (key datastore.Key, ok bool) {
 }
 
 func (fs *Datastore) makePrefixDir(dir string) error {
-	if err := os.Mkdir(dir, 0777); err != nil {
-		// EEXIST is safe to ignore here, that just means the prefix
-		// directory already existed.
-		if !os.IsExist(err) {
-			return err
-		}
+	if err := fs.makePrefixDirNoSync(dir); err != nil {
+		return err
 	}
 
 	// In theory, if we create a new prefix dir and add a file to
@@ -86,6 +82,17 @@ func (fs *Datastore) makePrefixDir(dir string) error {
 	return nil
 }
 
+func (fs *Datastore) makePrefixDirNoSync(dir string) error {
+	if err := os.Mkdir(dir, 0777); err != nil {
+		// EEXIST is safe to ignore here, that just means the prefix
+		// directory already existed.
+		if !os.IsExist(err) {
+			return err
+		}
+	}
+	return nil
+}
+
 func (fs *Datastore) Put(key datastore.Key, value interface{}) error {
 	val, ok := value.([]byte)
 	if !ok {
@@ -137,6 +144,88 @@ func (fs *Datastore) Put(key datastore.Key, value interface{}) error {
 	return nil
 }
 
+func (fs *Datastore) putMany(data map[datastore.Key]interface{}) error {
+	var dirsToSync []string
+	files := make(map[*os.File]string)
+
+	for key, value := range data {
+		val, ok := value.([]byte)
+		if !ok {
+			return datastore.ErrInvalidType
+		}
+		dir, path := fs.encode(key)
+		if err := fs.makePrefixDirNoSync(dir); err != nil {
+			return err
+		}
+		dirsToSync = append(dirsToSync, dir)
+
+		tmp, err := ioutil.TempFile(dir, "put-")
+		if err != nil {
+			return err
+		}
+
+		if _, err := tmp.Write(val); err != nil {
+			return err
+		}
+
+		files[tmp] = path
+	}
+
+	ops := make(map[*os.File]int)
+
+	defer func() {
+		for fi, _ := range files {
+			val, _ := ops[fi]
+			switch val {
+			case 0:
+				_ = fi.Close()
+				fallthrough
+			case 1:
+				_ = os.Remove(fi.Name())
+			}
+		}
+	}()
+
+	// Now we sync everything
+	// sync and close files
+	for fi, _ := range files {
+		if err := fi.Sync(); err != nil {
+			return err
+		}
+
+		if err := fi.Close(); err != nil {
+			return err
+		}
+
+		// signify closed
+		ops[fi] = 1
+	}
+
+	// move files to their proper places
+	for fi, path := range files {
+		if err := osrename.Rename(fi.Name(), path); err != nil {
+			return err
+		}
+
+		// signify removed
+		ops[fi] = 2
+	}
+
+	// now sync the dirs for those files
+	for _, dir := range dirsToSync {
+		if err := syncDir(dir); err != nil {
+			return err
+		}
+	}
+
+	// sync top flatfs dir
+	if err := syncDir(fs.path); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (fs *Datastore) Get(key datastore.Key) (value interface{}, err error) {
 	_, path := fs.encode(key)
 	data, err := ioutil.ReadFile(path)
@@ -234,6 +323,51 @@ func (fs *Datastore) enumerateKeys(fi os.FileInfo, res []query.Entry) ([]query.E
 	return res, nil
 }
 
+type flatfsTransaction struct {
+	puts    map[datastore.Key]interface{}
+	gets    map[datastore.Key]datastore.GetCallback
+	deletes map[datastore.Key]struct{}
+
+	ds *Datastore
+}
+
+func (fs *Datastore) StartBatchOp() datastore.Transaction {
+	return &flatfsTransaction{
+		puts:    make(map[datastore.Key]interface{}),
+		gets:    make(map[datastore.Key]datastore.GetCallback),
+		deletes: make(map[datastore.Key]struct{}),
+		ds:      fs,
+	}
+}
+
+func (bt *flatfsTransaction) Put(key datastore.Key, val interface{}) error {
+	bt.puts[key] = val
+	return nil
+}
+
+func (bt *flatfsTransaction) Delete(key datastore.Key) error {
+	bt.deletes[key] = struct{}{}
+	return nil
+}
+
+func (bt *flatfsTransaction) Commit() error {
+	for k, cb := range bt.gets {
+		cb(bt.ds.Get(k))
+	}
+
+	if err := bt.ds.putMany(bt.puts); err != nil {
+		return err
+	}
+
+	for k, _ := range bt.deletes {
+		if err := bt.ds.Delete(k); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 var _ datastore.ThreadSafeDatastore = (*Datastore)(nil)
 
 func (*Datastore) IsThreadSafe() {}