core: ObjectStorage.Begin and TxObjectStorage

14 files changed, 180 insertions, 38 deletions

diff --git a/blame_test.go b/blame_test.go
index 61d505a..a8f7ae1 100644
--- a/blame_test.go
+++ b/blame_test.go
@@ -31,8 +31,8 @@ func (s *BlameCommon) SetUpSuite(c *C) {
 
 		stream := packfile.NewStream(bytes.NewReader(data))
 
-		d := packfile.NewDecoder(stream)
-		err = d.Decode(r.s.ObjectStorage())
+		d := packfile.NewDecoder(stream, r.s.ObjectStorage())
+		err = d.Decode()
 		c.Assert(err, IsNil)
 
 		c.Assert(f.Close(), IsNil)
diff --git a/common_test.go b/common_test.go
index 42b3310..f054f23 100644
--- a/common_test.go
+++ b/common_test.go
@@ -123,8 +123,8 @@ func unpackFixtures(c *C, fixtures ...[]packedFixture) map[string]*Repository {
 			memStream := bytes.NewReader(data)
 			r := packfile.NewStream(memStream)
 
-			d := packfile.NewDecoder(r)
-			err = d.Decode(repos[fixture.url].s.ObjectStorage())
+			d := packfile.NewDecoder(r, repos[fixture.url].s.ObjectStorage())
+			err = d.Decode()
 			c.Assert(err, IsNil, comment)
 
 			c.Assert(f.Close(), IsNil, comment)
diff --git a/core/storage.go b/core/storage.go
index 15d988a..09b2450 100644
--- a/core/storage.go
+++ b/core/storage.go
@@ -7,19 +7,28 @@ var ErrStop = errors.New("stop iter")
 
 // ObjectStorage generic storage of objects
 type ObjectStorage interface {
+	// NewObject returns a new Object, the real type of the object can be a
+	// custom implementation or the defaul one, MemoryObject
 	NewObject() Object
+	// Set save an object into the storage, the object shuld be create with
+	// the NewObject, method, and file if the type is not supported.
 	Set(Object) (Hash, error)
-	// Get an object by hash with the given ObjectType.
-	//
-	// Implementors should return (nil, core.ErrObjectNotFound) if an object
-	// doesn't exist with both the given hash and object type.
+	// Get an object by hash with the given ObjectType. Implementors should
+	// return (nil, ErrObjectNotFound) if an object doesn't exist with both the
+	// given hash and object type.
 	//
 	// Valid ObjectType values are CommitObject, BlobObject, TagObject, TreeObject
 	// and AnyObject.
 	//
 	// If AnyObject is given, the object must be looked up regardless of its type.
 	Get(ObjectType, Hash) (Object, error)
+	// Iter returns a custom ObjectIter over all the object on the storage.
+	//
+	// Valid ObjectType values are CommitObject, BlobObject, TagObject, TreeObject
+	// and AnyObject.
 	Iter(ObjectType) (ObjectIter, error)
+	// Begin starts a transaction.
+	Begin() TxObjectStorage
 }
 
 // ObjectIter is a generic closable interface for iterating over objects.
@@ -29,6 +38,14 @@ type ObjectIter interface {
 	Close()
 }
 
+// TxObjectStorage is an in-progress storage transaction.
+// A transaction must end with a call to Commit or Rollback.
+type TxObjectStorage interface {
+	Set(Object) (Hash, error)
+	Commit() error
+	Rollback() error
+}
+
 // ReferenceStorage generic storage of references
 type ReferenceStorage interface {
 	Set(*Reference) error
diff --git a/examples/storage/aerospike/storage.go b/examples/storage/aerospike/storage.go
index e77448a..b57d932 100644
--- a/examples/storage/aerospike/storage.go
+++ b/examples/storage/aerospike/storage.go
@@ -131,6 +131,26 @@ func (s *ObjectStorage) buildKey(h core.Hash, t core.ObjectType) (*driver.Key, e
 	return driver.NewKey(s.ns, t.String(), fmt.Sprintf("%s|%s", s.url, h.String()))
 }
 
+func (s *ObjectStorage) Begin() core.TxObjectStorage {
+	return &TxObjectStorage{Storage: s}
+}
+
+type TxObjectStorage struct {
+	Storage *ObjectStorage
+}
+
+func (tx *TxObjectStorage) Set(obj core.Object) (core.Hash, error) {
+	return tx.Storage.Set(obj)
+}
+
+func (tx *TxObjectStorage) Commit() error {
+	return nil
+}
+
+func (tx *TxObjectStorage) Rollback() error {
+	return nil
+}
+
 type ObjectIter struct {
 	t  core.ObjectType
 	ch chan *driver.Record
diff --git a/formats/packfile/decoder.go b/formats/packfile/decoder.go
index 6d0cd8b..5b5763c 100644
--- a/formats/packfile/decoder.go
+++ b/formats/packfile/decoder.go
@@ -42,25 +42,33 @@ type Decoder struct {
 }
 
 // NewDecoder returns a new Decoder that reads from r.
-func NewDecoder(r ReadRecaller) *Decoder {
+func NewDecoder(r ReadRecaller, s core.ObjectStorage) *Decoder {
 	return &Decoder{
 		p: NewParser(r),
+		s: s,
 	}
 }
 
 // Decode reads a packfile and stores it in the value pointed to by s.
-func (d *Decoder) Decode(s core.ObjectStorage) error {
-	d.s = s
-
+func (d *Decoder) Decode() error {
 	count, err := d.p.ReadHeader()
 	if err != nil {
 		return err
 	}
 
-	return d.readObjects(count)
+	tx := d.s.Begin()
+	if err := d.readObjects(tx, count); err != nil {
+		if err := tx.Rollback(); err != nil {
+			return nil
+		}
+
+		return err
+	}
+
+	return tx.Commit()
 }
 
-func (d *Decoder) readObjects(count uint32) error {
+func (d *Decoder) readObjects(tx core.TxObjectStorage, count uint32) error {
 	// This code has 50-80 µs of overhead per object not counting zlib inflation.
 	// Together with zlib inflation, it's 400-410 µs for small objects.
 	// That's 1 sec for ~2450 objects, ~4.20 MB, or ~250 ms per MB,
@@ -85,7 +93,7 @@ func (d *Decoder) readObjects(count uint32) error {
 			return err
 		}
 
-		_, err = d.s.Set(obj)
+		_, err = tx.Set(obj)
 		if err == io.EOF {
 			break
 		}
diff --git a/formats/packfile/decoder_test.go b/formats/packfile/decoder_test.go
index fe25ee4..8c73b4e 100644
--- a/formats/packfile/decoder_test.go
+++ b/formats/packfile/decoder_test.go
@@ -28,10 +28,10 @@ func (s *ReaderSuite) TestReadPackfile(c *C) {
 	data, _ := base64.StdEncoding.DecodeString(packFileWithEmptyObjects)
 	f := bytes.NewReader(data)
 	r := NewStream(f)
-	d := NewDecoder(r)
-
 	sto := memory.NewStorage()
-	err := d.Decode(sto.ObjectStorage())
+	d := NewDecoder(r, sto.ObjectStorage())
+
+	err := d.Decode()
 	c.Assert(err, IsNil)
 
 	AssertObjects(c, sto, []string{
@@ -61,10 +61,10 @@ func (s *ReaderSuite) testReadPackfileGitFixture(c *C, file string, format Forma
 	f, err := os.Open(file)
 	c.Assert(err, IsNil)
 	r := NewSeekable(f)
-	d := NewDecoder(r)
-
 	sto := memory.NewStorage()
-	err = d.Decode(sto.ObjectStorage())
+	d := NewDecoder(r, sto.ObjectStorage())
+
+	err = d.Decode()
 	c.Assert(err, IsNil)
 
 	AssertObjects(c, sto, []string{
@@ -170,10 +170,10 @@ func readFromFile(c *C, file string, format Format) *memory.ObjectStorage {
 	f, err := os.Open(file)
 	c.Assert(err, IsNil)
 	r := NewSeekable(f)
-	d := NewDecoder(r)
-
 	sto := memory.NewStorage()
-	err = d.Decode(sto.ObjectStorage())
+	d := NewDecoder(r, sto.ObjectStorage())
+
+	err = d.Decode()
 	c.Assert(err, IsNil)
 
 	return sto.ObjectStorage().(*memory.ObjectStorage)
diff --git a/references_test.go b/references_test.go
index 216fadf..3907b25 100644
--- a/references_test.go
+++ b/references_test.go
@@ -27,8 +27,8 @@ func (s *ReferencesSuite) SetUpSuite(c *C) {
 		defer f.Close()
 		c.Assert(err, IsNil)
 		r := packfile.NewSeekable(f)
-		d := packfile.NewDecoder(r)
-		err = d.Decode(s.repos[fix.url].s.ObjectStorage())
+		d := packfile.NewDecoder(r, s.repos[fix.url].s.ObjectStorage())
+		err = d.Decode()
 		c.Assert(err, IsNil)
 	}
 }
diff --git a/remote.go b/remote.go
index f388d87..2605d39 100644
--- a/remote.go
+++ b/remote.go
@@ -163,8 +163,8 @@ func (r *Remote) buildRequest(
 func (r *Remote) updateObjectStorage(reader io.Reader) error {
 	stream := packfile.NewStream(reader)
 
-	d := packfile.NewDecoder(stream)
-	return d.Decode(r.s.ObjectStorage())
+	d := packfile.NewDecoder(stream, r.s.ObjectStorage())
+	return d.Decode()
 }
 
 func (r *Remote) updateLocalReferenceStorage(specs []config.RefSpec, refs []*core.Reference) error {
diff --git a/storage/filesystem/object.go b/storage/filesystem/object.go
index 3888fd8..0696c70 100644
--- a/storage/filesystem/object.go
+++ b/storage/filesystem/object.go
@@ -217,3 +217,20 @@ func buildIndexFromIdxfile(fs fs.FS, path string) (index.Index, error) {
 
 	return index.NewFromIdx(f)
 }
+func (o *ObjectStorage) Begin() core.TxObjectStorage {
+	return &TxObjectStorage{}
+}
+
+type TxObjectStorage struct{}
+
+func (tx *TxObjectStorage) Set(obj core.Object) (core.Hash, error) {
+	return core.ZeroHash, fmt.Errorf("not implemented yet")
+}
+
+func (tx *TxObjectStorage) Commit() error {
+	return fmt.Errorf("not implemented yet")
+}
+
+func (tx *TxObjectStorage) Rollback() error {
+	return fmt.Errorf("not implemented yet")
+}
diff --git a/storage/filesystem/object_test.go b/storage/filesystem/object_test.go
index 361ae32..5752f79 100644
--- a/storage/filesystem/object_test.go
+++ b/storage/filesystem/object_test.go
@@ -124,9 +124,9 @@ func memStorageFromGitDir(fs fs.FS, path string) (core.ObjectStorage, error) {
 
 	sto := memory.NewStorage()
 	r := packfile.NewStream(f)
-	d := packfile.NewDecoder(r)
+	d := packfile.NewDecoder(r, sto.ObjectStorage())
 
-	err = d.Decode(sto.ObjectStorage())
+	err = d.Decode()
 	if err != nil {
 		return nil, err
 	}
@@ -267,8 +267,8 @@ func memStorageFromDirPath(fs fs.FS, path string) (core.ObjectStorage, error) {
 	}
 
 	r := packfile.NewStream(f)
-	d := packfile.NewDecoder(r)
-	err = d.Decode(sto.ObjectStorage())
+	d := packfile.NewDecoder(r, sto.ObjectStorage())
+	err = d.Decode()
 	if err != nil {
 		return nil, err
 	}
diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index c8a152d..01d5fa5 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -166,6 +166,41 @@ func flattenObjectMap(m map[core.Hash]core.Object) []core.Object {
 	return objects
 }
 
+func (o *ObjectStorage) Begin() core.TxObjectStorage {
+	return &TxObjectStorage{
+		Storage: o,
+		Objects: make(map[core.Hash]core.Object, 0),
+	}
+}
+
+type TxObjectStorage struct {
+	Storage *ObjectStorage
+	Objects map[core.Hash]core.Object
+}
+
+func (tx *TxObjectStorage) Set(obj core.Object) (core.Hash, error) {
+	h := obj.Hash()
+	tx.Objects[h] = obj
+
+	return h, nil
+}
+
+func (tx *TxObjectStorage) Commit() error {
+	for h, obj := range tx.Objects {
+		delete(tx.Objects, h)
+		if _, err := tx.Storage.Set(obj); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (tx *TxObjectStorage) Rollback() error {
+	tx.Objects = make(map[core.Hash]core.Object, 0)
+	return nil
+}
+
 type ReferenceStorage map[core.ReferenceName]*core.Reference
 
 // Set stores a reference.
diff --git a/storage/memory/storage_test.go b/storage/memory/storage_test.go
index 61a810e..2c90738 100644
--- a/storage/memory/storage_test.go
+++ b/storage/memory/storage_test.go
@@ -6,19 +6,19 @@ import (
 
 	. "gopkg.in/check.v1"
 	"gopkg.in/src-d/go-git.v4/core"
-	. "gopkg.in/src-d/go-git.v4/storage/test"
+	"gopkg.in/src-d/go-git.v4/storage/test"
 )
 
 func Test(t *testing.T) { TestingT(t) }
 
 type StorageSuite struct {
-	BaseStorageSuite
+	test.BaseStorageSuite
 }
 
 var _ = Suite(&StorageSuite{})
 
-func (s *StorageSuite) SetUpSuite(c *C) {
-	s.BaseStorageSuite = NewBaseStorageSuite(NewStorage().ObjectStorage())
+func (s *StorageSuite) SetUpTest(c *C) {
+	s.BaseStorageSuite = test.NewBaseStorageSuite(NewStorage().ObjectStorage())
 }
 
 func (s *StorageSuite) TestStorageObjectStorage(c *C) {
diff --git a/storage/test/storage_suite.go b/storage/test/storage_suite.go
index c68a089..b4ba306 100644
--- a/storage/test/storage_suite.go
+++ b/storage/test/storage_suite.go
@@ -114,3 +114,48 @@ func (s *BaseStorageSuite) TestObjectStorageIter(c *C) {
 		c.Assert(found, Equals, true, Commentf("Object of type %s not found", to.Type.String()))
 	}
 }
+
+func (s *BaseStorageSuite) TestTxObjectStorageSetAndCommit(c *C) {
+	tx := s.ObjectStorage.Begin()
+	for _, o := range s.testObjects {
+		h, err := tx.Set(o.Object)
+		c.Assert(err, IsNil)
+		c.Assert(h.String(), Equals, o.Hash)
+	}
+
+	iter, err := s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+	_, err = iter.Next()
+	c.Assert(err, Equals, io.EOF)
+
+	err = tx.Commit()
+	c.Assert(err, IsNil)
+
+	iter, err = s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+
+	var count int
+	iter.ForEach(func(o core.Object) error {
+		count++
+		return nil
+	})
+
+	c.Assert(count, Equals, 4)
+}
+
+func (s *BaseStorageSuite) TestTxObjectStorageSetAndRollback(c *C) {
+	tx := s.ObjectStorage.Begin()
+	for _, o := range s.testObjects {
+		h, err := tx.Set(o.Object)
+		c.Assert(err, IsNil)
+		c.Assert(h.String(), Equals, o.Hash)
+	}
+
+	err := tx.Rollback()
+	c.Assert(err, IsNil)
+
+	iter, err := s.ObjectStorage.Iter(core.AnyObject)
+	c.Assert(err, IsNil)
+	_, err = iter.Next()
+	c.Assert(err, Equals, io.EOF)
+}
diff --git a/tree_diff_test.go b/tree_diff_test.go
index b7efaa1..e177c26 100644
--- a/tree_diff_test.go
+++ b/tree_diff_test.go
@@ -49,8 +49,8 @@ func (s *DiffTreeSuite) SetUpSuite(c *C) {
 		c.Assert(err, IsNil)
 
 		r := packfile.NewSeekable(f)
-		d := packfile.NewDecoder(r)
-		err = d.Decode(s.repos[fixRepo.url].s.ObjectStorage())
+		d := packfile.NewDecoder(r, s.repos[fixRepo.url].s.ObjectStorage())
+		err = d.Decode()
 		c.Assert(err, IsNil)
 
 		c.Assert(f.Close(), IsNil)