Move the upload flag to the database, and use the database stored image size rather than going to s3.

data/database.py

@@ -220,7 +220,8 @@ class ImageStorage(BaseModel):
   created = DateTimeField(null=True)
   comment = TextField(null=True)
   command = TextField(null=True)
-  image_size = BigIntegerField(null=True)  
+  image_size = BigIntegerField(null=True)
+  uploading = BooleanField(default=True, null=True)
 
 
 class Image(BaseModel):

data/model/legacy.py

@@ -817,7 +817,7 @@ def get_repository(namespace_name, repository_name):
 
 def get_repo_image(namespace_name, repository_name, image_id):
   query = (Image
-    .select()
+    .select(Image, ImageStorage)
     .join(Repository)
     .switch(Image)
     .join(ImageStorage, JOIN_LEFT_OUTER)

endpoints/registry.py

@@ -40,16 +40,35 @@ class SocketReader(object):
     return buf
 
 
+def image_is_uploading(namespace, repository, image_id, repo_image):
+  if repo_image and repo_image.storage and repo_image.storage.uploading is not None:
+    return repo_image.storage.uploading
+
+  logger.warning('Setting legacy upload flag')
+  uuid = repo_image and repo_image.storage and repo_image.storage.uuid
+  mark_path = store.image_mark_path(namespace, repository, image_id, uuid)
+  return store.exists(mark_path)
+
+
+def mark_upload_complete(namespace, repository, image_id, repo_image):
+  if repo_image and repo_image.storage and repo_image.storage.uploading is not None:
+    repo_image.storage.uploading = False
+    repo_image.storage.save()
+  else:
+    logger.warning('Removing legacy upload flag')
+    uuid = repo_image and repo_image.storage and repo_image.storage.uuid
+    mark_path = store.image_mark_path(namespace, repository, image_id, uuid)
+    if store.exists(mark_path):
+      store.remove(mark_path)
+
+
 def require_completion(f):
   """This make sure that the image push correctly finished."""
   @wraps(f)
   def wrapper(namespace, repository, *args, **kwargs):
     image_id = kwargs['image_id']
     repo_image = model.get_repo_image(namespace, repository, image_id)
-    uuid = repo_image and repo_image.storage and repo_image.storage.uuid
-
-    if store.exists(store.image_mark_path(namespace, repository, image_id,
-                                          uuid)):
+    if image_is_uploading(namespace, repository, image_id, repo_image):
       abort(400, 'Image %(image_id)s is being uploaded, retry later',
             issue='upload-in-progress', image_id=kwargs['image_id'])
 
@@ -139,9 +158,9 @@ def put_image_layer(namespace, repository, image_id):
 
   profile.debug('Retrieving image path info')
   layer_path = store.image_layer_path(namespace, repository, image_id, uuid)
-  mark_path = store.image_mark_path(namespace, repository, image_id, uuid)
 
-  if store.exists(layer_path) and not store.exists(mark_path):
+  if (store.exists(layer_path) and not
+      image_is_uploading(namespace, repository, image_id, repo_image)):
     abort(409, 'Image already exists', issue='image-exists', image_id=image_id)
 
   profile.debug('Storing layer data')
@@ -192,7 +211,7 @@ def put_image_layer(namespace, repository, image_id):
           issue='checksum-mismatch', image_id=image_id)
 
   # Checksum is ok, we remove the marker
-  store.remove(mark_path)
+  mark_upload_complete(namespace, repository, image_id, repo_image)
 
   # The layer is ready for download, send a job to the work queue to
   # process it.
@@ -234,8 +253,7 @@ def put_image_checksum(namespace, repository, image_id):
     abort(404, 'Image not found: %(image_id)s', issue='unknown-image', image_id=image_id)
 
   profile.debug('Marking image path')
-  mark_path = store.image_mark_path(namespace, repository, image_id, uuid)
-  if not store.exists(mark_path):
+  if not image_is_uploading(namespace, repository, image_id, repo_image):
     abort(409, 'Cannot set checksum for image %(image_id)s',
           issue='image-write-error', image_id=image_id)
 
@@ -252,7 +270,7 @@ def put_image_checksum(namespace, repository, image_id):
           issue='checksum-mismatch', image_id=image_id)
 
   # Checksum is ok, we remove the marker
-  store.remove(mark_path)
+  mark_upload_complete(namespace, repository, image_id, repo_image)
 
   # The layer is ready for download, send a job to the work queue to
   # process it.
@@ -291,8 +309,7 @@ def get_image_json(namespace, repository, image_id, headers):
 
   profile.debug('Looking up repo layer size')
   try:
-    size = store.get_size(store.image_layer_path(namespace, repository,
-                                                 image_id, uuid))
+    size = repo_image.image_size or repo_image.storage.image_size
     headers['X-Docker-Size'] = str(size)
   except OSError:
     pass
@@ -429,13 +446,13 @@ def put_image_json(namespace, repository, image_id):
                                          parent_uuid))):
     abort(400, 'Image %(image_id)s depends on non existing parent image %(parent_id)s',
           issue='invalid-request', image_id=image_id, parent_id=parent_id)
-    
+
   profile.debug('Looking up image storage paths')
   json_path = store.image_json_path(namespace, repository, image_id, uuid)
-  mark_path = store.image_mark_path(namespace, repository, image_id, uuid)
 
   profile.debug('Checking if image already exists')
-  if store.exists(json_path) and not store.exists(mark_path):
+  if (store.exists(json_path) and not
+      image_is_uploading(namespace, repository, image_id, repo_image)):
     abort(409, 'Image already exists', issue='image-exists', image_id=image_id)
 
   # If we reach that point, it means that this is a new image or a retry
@@ -449,9 +466,6 @@ def put_image_json(namespace, repository, image_id):
                            data.get('created'), data.get('comment'), command,
                            parent_image)
 
-  profile.debug('Putting mark path')
-  store.put_content(mark_path, 'true')
-
   profile.debug('Putting json path')
   store.put_content(json_path, request.data)
 

storage/basestorage.py

@@ -92,6 +92,3 @@ class BaseStorage(object):
 
   def remove(self, path):
     raise NotImplementedError
-
-  def get_size(self, path):
-    raise NotImplementedError

storage/local.py

@@ -80,7 +80,3 @@ class LocalStorage(BaseStorage):
       os.remove(path)
     except OSError:
       pass
-
-  def get_size(self, path):
-    path = self._init_path(path)
-    return os.path.getsize(path)

storage/s3.py

@@ -171,12 +171,3 @@ class S3Storage(BaseStorage):
       path += '/'
     for key in self._s3_bucket.list(prefix=path):
       key.delete()
-
-  def get_size(self, path):
-    self._initialize_s3()
-    path = self._init_path(path)
-    # Lookup does a HEAD HTTP Request on the object
-    key = self._s3_bucket.lookup(path)
-    if not key:
-      raise OSError('No such key: \'{0}\''.format(path))
-    return key.size