5 files changed, 205 insertions, 0 deletions

diff --git a/plumbing/storer/reference.go b/plumbing/storer/reference.go
index ae80a39..5e85a3b 100644
--- a/plumbing/storer/reference.go
+++ b/plumbing/storer/reference.go
@@ -24,6 +24,8 @@ type ReferenceStorer interface {
 	Reference(plumbing.ReferenceName) (*plumbing.Reference, error)
 	IterReferences() (ReferenceIter, error)
 	RemoveReference(plumbing.ReferenceName) error
+	CountLooseRefs() (int, error)
+	PackRefs() error
 }
 
 // ReferenceIter is a generic closable interface for iterating over references.
diff --git a/storage/filesystem/internal/dotgit/dotgit.go b/storage/filesystem/internal/dotgit/dotgit.go
index 1cb97bd..29e2525 100644
--- a/storage/filesystem/internal/dotgit/dotgit.go
+++ b/storage/filesystem/internal/dotgit/dotgit.go
@@ -576,6 +576,121 @@ func (d *DotGit) readReferenceFile(path, name string) (ref *plumbing.Reference,
 	return d.readReferenceFrom(f, name)
 }
 
+func (d *DotGit) CountLooseRefs() (int, error) {
+	var refs []*plumbing.Reference
+	var seen = make(map[plumbing.ReferenceName]bool)
+	if err := d.addRefsFromRefDir(&refs, seen); err != nil {
+		return 0, err
+	}
+
+	return len(refs), nil
+}
+
+// PackRefs packs all loose refs into the packed-refs file.
+//
+// This implementation only works under the assumption that the view
+// of the file system won't be updated during this operation, which is
+// true for kbfsgit after the Lock() operation is complete (and before
+// the Unlock()/Close() of the locked file).  If another process
+// concurrently updates one of the loose refs we delete, then KBFS
+// conflict resolution would just end up ignoring our delete.  Also
+// note that deleting a ref requires locking packed-refs, so a ref
+// deleted by the user shouldn't be revived by ref-packing.
+//
+// The strategy would not work on a general file system though,
+// without locking each loose reference and checking it again before
+// deleting the file, because otherwise an updated reference could
+// sneak in and then be deleted by the packed-refs process.
+// Alternatively, every ref update could also lock packed-refs, so
+// only one lock is required during ref-packing.  But that would
+// worsen performance in the common case.
+//
+// TODO: before trying to get this merged upstream, move it into a
+// custom kbfsgit Storer implementation, and rewrite this function to
+// work correctly on a general filesystem.
+func (d *DotGit) PackRefs() (err error) {
+	// Lock packed-refs, and create it if it doesn't exist yet.
+	f, err := d.fs.Open(packedRefsPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			f, err = d.fs.Create(packedRefsPath)
+			if err != nil {
+				return err
+			}
+		} else {
+			return err
+		}
+	}
+	defer ioutil.CheckClose(f, &err)
+
+	err = f.Lock()
+	if err != nil {
+		return err
+	}
+
+	// Gather all refs using addRefsFromRefDir and addRefsFromPackedRefs.
+	var refs []*plumbing.Reference
+	var seen = make(map[plumbing.ReferenceName]bool)
+	if err := d.addRefsFromRefDir(&refs, seen); err != nil {
+		return err
+	}
+	if len(refs) == 0 {
+		// Nothing to do!
+		return nil
+	}
+	numLooseRefs := len(refs)
+	if err := d.addRefsFromPackedRefs(&refs, seen); err != nil {
+		return err
+	}
+
+	// Write them all to a new temp packed-refs file.
+	tmp, err := d.fs.TempFile("", tmpPackedRefsPrefix)
+	if err != nil {
+		return err
+	}
+	doCloseTmp := true
+	defer func() {
+		if doCloseTmp {
+			ioutil.CheckClose(tmp, &err)
+		}
+	}()
+	for _, ref := range refs {
+		_, err := tmp.Write([]byte(ref.String() + "\n"))
+		if err != nil {
+			return err
+		}
+	}
+
+	// Rename the temp packed-refs file.
+	doCloseTmp = false
+	if err := tmp.Close(); err != nil {
+		return err
+	}
+	err = d.fs.Rename(tmp.Name(), packedRefsPath)
+	if err != nil {
+		return err
+	}
+
+	// Delete all the loose refs, while still holding the packed-refs
+	// lock.
+	for _, ref := range refs[:numLooseRefs] {
+		path := d.fs.Join(".", ref.Name().String())
+		err = d.fs.Remove(path)
+		if err != nil && !os.IsNotExist(err) {
+			return err
+		}
+	}
+
+	// Update packed-refs cache.
+	d.cachedPackedRefs = make(refCache)
+	for _, ref := range refs {
+		d.cachedPackedRefs[ref.Name()] = ref
+	}
+	d.packedRefsLastMod = time.Now()
+
+	return nil
+}
+
 // Module return a billy.Filesystem poiting to the module folder
 func (d *DotGit) Module(name string) (billy.Filesystem, error) {
 	return d.fs.Chroot(d.fs.Join(modulePath, name))
diff --git a/storage/filesystem/internal/dotgit/dotgit_test.go b/storage/filesystem/internal/dotgit/dotgit_test.go
index 446a204..0ed9ec5 100644
--- a/storage/filesystem/internal/dotgit/dotgit_test.go
+++ b/storage/filesystem/internal/dotgit/dotgit_test.go
@@ -543,3 +543,75 @@ func (s *SuiteDotGit) TestSubmodules(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(strings.HasSuffix(m.Root(), m.Join(".git", "modules", "basic")), Equals, true)
 }
+
+func (s *SuiteDotGit) TestPackRefs(c *C) {
+	tmp, err := ioutil.TempDir("", "dot-git")
+	c.Assert(err, IsNil)
+	defer os.RemoveAll(tmp)
+
+	fs := osfs.New(tmp)
+	dir := New(fs)
+
+	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+		"refs/heads/foo",
+		"e8d3ffab552895c19b9fcf7aa264d277cde33881",
+	), nil)
+	c.Assert(err, IsNil)
+	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+		"refs/heads/bar",
+		"a8d3ffab552895c19b9fcf7aa264d277cde33881",
+	), nil)
+	c.Assert(err, IsNil)
+
+	refs, err := dir.Refs()
+	c.Assert(err, IsNil)
+	c.Assert(refs, HasLen, 2)
+	looseCount, err := dir.CountLooseRefs()
+	c.Assert(err, IsNil)
+	c.Assert(looseCount, Equals, 2)
+
+	err = dir.PackRefs()
+	c.Assert(err, IsNil)
+
+	// Make sure the refs are still there, but no longer loose.
+	refs, err = dir.Refs()
+	c.Assert(err, IsNil)
+	c.Assert(refs, HasLen, 2)
+	looseCount, err = dir.CountLooseRefs()
+	c.Assert(err, IsNil)
+	c.Assert(looseCount, Equals, 0)
+
+	ref, err := dir.Ref("refs/heads/foo")
+	c.Assert(err, IsNil)
+	c.Assert(ref, NotNil)
+	c.Assert(ref.Hash().String(), Equals, "e8d3ffab552895c19b9fcf7aa264d277cde33881")
+	ref, err = dir.Ref("refs/heads/bar")
+	c.Assert(err, IsNil)
+	c.Assert(ref, NotNil)
+	c.Assert(ref.Hash().String(), Equals, "a8d3ffab552895c19b9fcf7aa264d277cde33881")
+
+	// Now update one of them, re-pack, and check again.
+	err = dir.SetRef(plumbing.NewReferenceFromStrings(
+		"refs/heads/foo",
+		"b8d3ffab552895c19b9fcf7aa264d277cde33881",
+	), nil)
+	c.Assert(err, IsNil)
+	looseCount, err = dir.CountLooseRefs()
+	c.Assert(err, IsNil)
+	c.Assert(looseCount, Equals, 1)
+	err = dir.PackRefs()
+	c.Assert(err, IsNil)
+
+	// Make sure the refs are still there, but no longer loose.
+	refs, err = dir.Refs()
+	c.Assert(err, IsNil)
+	c.Assert(refs, HasLen, 2)
+	looseCount, err = dir.CountLooseRefs()
+	c.Assert(err, IsNil)
+	c.Assert(looseCount, Equals, 0)
+
+	ref, err = dir.Ref("refs/heads/foo")
+	c.Assert(err, IsNil)
+	c.Assert(ref, NotNil)
+	c.Assert(ref.Hash().String(), Equals, "b8d3ffab552895c19b9fcf7aa264d277cde33881")
+}
diff --git a/storage/filesystem/reference.go b/storage/filesystem/reference.go
index 54cdf56..7313f05 100644
--- a/storage/filesystem/reference.go
+++ b/storage/filesystem/reference.go
@@ -34,3 +34,11 @@ func (r *ReferenceStorage) IterReferences() (storer.ReferenceIter, error) {
 func (r *ReferenceStorage) RemoveReference(n plumbing.ReferenceName) error {
 	return r.dir.RemoveRef(n)
 }
+
+func (r *ReferenceStorage) CountLooseRefs() (int, error) {
+	return r.dir.CountLooseRefs()
+}
+
+func (r *ReferenceStorage) PackRefs() error {
+	return r.dir.PackRefs()
+}
diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index 927ec41..3d4e84a 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -236,6 +236,14 @@ func (r ReferenceStorage) IterReferences() (storer.ReferenceIter, error) {
 	return storer.NewReferenceSliceIter(refs), nil
 }
 
+func (r ReferenceStorage) CountLooseRefs() (int, error) {
+	return len(r), nil
+}
+
+func (r ReferenceStorage) PackRefs() error {
+	return nil
+}
+
 func (r ReferenceStorage) RemoveReference(n plumbing.ReferenceName) error {
 	delete(r, n)
 	return nil