From 02cb7f45fa6f9a8899286431fa68bf1a9d0194d8 Mon Sep 17 00:00:00 2001
From: Solomon Hykes <solomon@dotcloud.com>
Date: Tue, 26 Nov 2013 04:48:34 +0000
Subject: [PATCH] Fix a race condition in TestInterruptedRegister

---
 graph.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/graph.go b/graph.go
index 82768d07dc..04d77b9146 100644
--- a/graph.go
+++ b/graph.go
@@ -131,7 +131,15 @@ func (graph *Graph) Create(layerData archive.Archive, container *Container, comm
 
 // Register imports a pre-existing image into the graph.
 // FIXME: pass img as first argument
-func (graph *Graph) Register(jsonData []byte, layerData archive.Archive, img *Image) error {
+func (graph *Graph) Register(jsonData []byte, layerData archive.Archive, img *Image) (err error) {
+	defer func() {
+		// If any error occurs, remove the new dir from the driver.
+		// Don't check for errors since the dir might not have been created.
+		// FIXME: this leaves a possible race condition.
+		if err != nil {
+			graph.driver.Remove(img.ID)
+		}
+	}()
 	if err := ValidateID(img.ID); err != nil {
 		return err
 	}
@@ -147,6 +155,12 @@ func (graph *Graph) Register(jsonData []byte, layerData archive.Archive, img *Im
 		return err
 	}
 
+	// If the driver has this ID but the graph doesn't, remove it from the driver to start fresh.
+	// (the graph is the source of truth).
+	// Ignore errors, since we don't know if the driver correctly returns ErrNotExist.
+	// (FIXME: make that mandatory for drivers).
+	graph.driver.Remove(img.ID)
+
 	tmp, err := graph.Mktemp("")
 	defer os.RemoveAll(tmp)
 	if err != nil {