Use optionally locking when updating refs

Author: taruti

plumbing/storer/reference.go

@@ -16,6 +16,7 @@ var ErrMaxResolveRecursion = errors.New("max. recursion level reached")
 // ReferenceStorer is a generic storage of references.
 type ReferenceStorer interface {
 	SetReference(*plumbing.Reference) error
+	CheckAndSetReference(new, old *plumbing.Reference) error
 	Reference(plumbing.ReferenceName) (*plumbing.Reference, error)
 	IterReferences() (ReferenceIter, error)
 	RemoveReference(plumbing.ReferenceName) error

remote.go

@@ -736,7 +736,7 @@ func (r *Remote) updateLocalReferenceStorage(
 				}
 			}
 
-			refUpdated, err := updateReferenceStorerIfNeeded(r.s, new)
+			refUpdated, err := checkAndUpdateReferenceStorerIfNeeded(r.s, new, old)
 			if err != nil {
 				return updated, err
 			}

repository.go

@@ -625,17 +625,17 @@ func (r *Repository) calculateRemoteHeadReference(spec []config.RefSpec,
 	return refs
 }
 
-func updateReferenceStorerIfNeeded(
-	s storer.ReferenceStorer, r *plumbing.Reference) (updated bool, err error) {
-
+func checkAndUpdateReferenceStorerIfNeeded(
+	s storer.ReferenceStorer, r, old *plumbing.Reference) (
+	updated bool, err error) {
 	p, err := s.Reference(r.Name())
 	if err != nil && err != plumbing.ErrReferenceNotFound {
 		return false, err
 	}
 
 	// we use the string method to compare references, is the easiest way
 	if err == plumbing.ErrReferenceNotFound || r.String() != p.String() {
-		if err := s.SetReference(r); err != nil {
+		if err := s.CheckAndSetReference(r, old); err != nil {
 			return false, err
 		}
 
@@ -645,6 +645,11 @@ func updateReferenceStorerIfNeeded(
 	return false, nil
 }
 
+func updateReferenceStorerIfNeeded(
+	s storer.ReferenceStorer, r *plumbing.Reference) (updated bool, err error) {
+	return checkAndUpdateReferenceStorerIfNeeded(s, r, nil)
+}
+
 // Fetch fetches references along with the objects necessary to complete
 // their histories, from the remote named as FetchOptions.RemoteName.
 //

storage/filesystem/internal/dotgit/dotgit.go

@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"errors"
 	"fmt"
+	"io"
 	stdioutil "io/ioutil"
 	"os"
 	"strings"
@@ -242,7 +243,39 @@ func (d *DotGit) Object(h plumbing.Hash) (billy.File, error) {
 	return d.fs.Open(file)
 }
 
-func (d *DotGit) SetRef(r *plumbing.Reference) error {
+func (d *DotGit) readReferenceFrom(rd io.Reader, name string) (ref *plumbing.Reference, err error) {
+	b, err := stdioutil.ReadAll(rd)
+	if err != nil {
+		return nil, err
+	}
+
+	line := strings.TrimSpace(string(b))
+	return plumbing.NewReferenceFromStrings(name, line), nil
+}
+
+func (d *DotGit) checkReferenceAndTruncate(f billy.File, old *plumbing.Reference) error {
+	if old == nil {
+		return nil
+	}
+	ref, err := d.readReferenceFrom(f, old.Name().String())
+	if err != nil {
+		return err
+	}
+	if ref.Hash() != old.Hash() {
+		return fmt.Errorf("reference has changed concurrently")
+	}
+	_, err = f.Seek(0, io.SeekStart)
+	if err != nil {
+		return err
+	}
+	err = f.Truncate(0)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (d *DotGit) SetRef(r, old *plumbing.Reference) error {
 	var content string
 	switch r.Type() {
 	case plumbing.SymbolicReference:
@@ -251,13 +284,30 @@ func (d *DotGit) SetRef(r *plumbing.Reference) error {
 		content = fmt.Sprintln(r.Hash().String())
 	}
 
-	f, err := d.fs.Create(r.Name().String())
+	// If we are not checking an old ref, just truncate the file.
+	mode := os.O_RDWR | os.O_CREATE
+	if old == nil {
+		mode |= os.O_TRUNC
+	}
+
+	f, err := d.fs.OpenFile(r.Name().String(), mode, 0666)
 	if err != nil {
 		return err
 	}
 
 	defer ioutil.CheckClose(f, &err)
 
+	err = f.Lock()
+	if err != nil {
+		return err
+	}
+
+	// this is a no-op to call even when old is nil.
+	err = d.checkReferenceAndTruncate(f, old)
+	if err != nil {
+		return err
+	}
+
 	_, err = f.Write([]byte(content))
 	return err
 }
@@ -512,13 +562,7 @@ func (d *DotGit) readReferenceFile(path, name string) (ref *plumbing.Reference,
 	}
 	defer ioutil.CheckClose(f, &err)
 
-	b, err := stdioutil.ReadAll(f)
-	if err != nil {
-		return nil, err
-	}
-
-	line := strings.TrimSpace(string(b))
-	return plumbing.NewReferenceFromStrings(name, line), nil
+	return d.readReferenceFrom(f, name)
 }
 
 // Module return a billy.Filesystem poiting to the module folder

storage/filesystem/reference.go

@@ -11,7 +11,11 @@ type ReferenceStorage struct {
 }
 
 func (r *ReferenceStorage) SetReference(ref *plumbing.Reference) error {
-	return r.dir.SetRef(ref)
+	return r.dir.SetRef(ref, nil)
+}
+
+func (r *ReferenceStorage) CheckAndSetReference(ref, old *plumbing.Reference) error {
+	return r.dir.SetRef(ref, old)
 }
 
 func (r *ReferenceStorage) Reference(n plumbing.ReferenceName) (*plumbing.Reference, error) {

storage/memory/storage.go

@@ -12,6 +12,7 @@ import (
 )
 
 var ErrUnsupportedObjectType = fmt.Errorf("unsupported object type")
+var ErrRefHasChanged = fmt.Errorf("reference has changed concurrently")
 
 // Storage is an implementation of git.Storer that stores data on memory, being
 // ephemeral. The use of this storage should be done in controlled envoriments,
@@ -202,6 +203,20 @@ func (r ReferenceStorage) SetReference(ref *plumbing.Reference) error {
 	return nil
 }
 
+func (r ReferenceStorage) CheckAndSetReference(ref, old *plumbing.Reference) error {
+	if ref != nil {
+		if old != nil {
+			tmp := r[ref.Name()]
+			if tmp != nil && tmp.Hash() != old.Hash() {
+				return ErrRefHasChanged
+			}
+		}
+		r[ref.Name()] = ref
+	}
+
+	return nil
+}
+
 func (r ReferenceStorage) Reference(n plumbing.ReferenceName) (*plumbing.Reference, error) {
 	ref, ok := r[n]
 	if !ok {