Merge remote-tracking branch 'upstream/master' into python-registry-v2

endpoints/v1/registry.py

@@ -1,12 +1,13 @@
 import logging
 import json
+import features
 
 from flask import make_response, request, session, Response, redirect, abort as flask_abort
 from functools import wraps
 from datetime import datetime
 from time import time
 
-from app import storage as store, image_diff_queue, app
+from app import storage as store, image_diff_queue, image_replication_queue, app
 from auth.auth import process_auth, extract_namespace_repo_from_session
 from auth.auth_context import get_authenticated_user, get_grant_user_context
 from digest import checksums
@@ -36,6 +37,30 @@ def set_uploading_flag(repo_image, is_image_uploading):
   repo_image.storage.save()
 
 
+def _finish_image(namespace, repository, repo_image):
+  # Checksum is ok, we remove the marker
+  set_uploading_flag(repo_image, False)
+
+  image_id = repo_image.docker_image_id
+
+  # The layer is ready for download, send a job to the work queue to
+  # process it.
+  logger.debug('Adding layer to diff queue')
+  repo = model.repository.get_repository(namespace, repository)
+  image_diff_queue.put([repo.namespace_user.username, repository, image_id], json.dumps({
+    'namespace_user_id': repo.namespace_user.id,
+    'repository': repository,
+    'image_id': image_id,
+  }))
+
+  # Send a job to the work queue to replicate the image layer.
+  if features.STORAGE_REPLICATION:
+    image_replication_queue.put([repo_image.storage.uuid], json.dumps({
+      'namespace_user_id': repo.namespace_user.id,
+      'storage_id': repo_image.storage.uuid,
+    }))
+
+
 def require_completion(f):
   """This make sure that the image push correctly finished."""
   @wraps(f)
@@ -210,7 +235,11 @@ def put_image_layer(namespace, repository, image_id):
 
   # Stream write the data to storage.
   with database.CloseForLongOperation(app.config):
-    store.stream_write(repo_image.storage.locations, layer_path, sr)
+    try:
+      store.stream_write(repo_image.storage.locations, layer_path, sr)
+    except IOError:
+      logger.exception('Exception when writing image data')
+      abort(520, 'Image %(image_id)s could not be written. Please try again.', image_id=image_id)
 
   # Append the computed checksum.
   csums = []
@@ -243,18 +272,8 @@ def put_image_layer(namespace, repository, image_id):
     abort(400, 'Checksum mismatch; ignoring the layer for image %(image_id)s',
           issue='checksum-mismatch', image_id=image_id)
 
-  # Checksum is ok, we remove the marker
-  set_uploading_flag(repo_image, False)
-
-  # The layer is ready for download, send a job to the work queue to
-  # process it.
-  logger.debug('Adding layer to diff queue')
-  repo = model.repository.get_repository(namespace, repository)
-  image_diff_queue.put([repo.namespace_user.username, repository, image_id], json.dumps({
-    'namespace_user_id': repo.namespace_user.id,
-    'repository': repository,
-    'image_id': image_id,
-  }))
+  # Mark the image as uploaded.
+  _finish_image(namespace, repository, repo_image)
 
   return make_response('true', 200)
 
@@ -316,18 +335,8 @@ def put_image_checksum(namespace, repository, image_id):
     abort(400, 'Checksum mismatch for image: %(image_id)s',
           issue='checksum-mismatch', image_id=image_id)
 
-  # Checksum is ok, we remove the marker
-  set_uploading_flag(repo_image, False)
-
-  # The layer is ready for download, send a job to the work queue to
-  # process it.
-  logger.debug('Adding layer to diff queue')
-  repo = model.repository.get_repository(namespace, repository)
-  image_diff_queue.put([repo.namespace_user.username, repository, image_id], json.dumps({
-    'namespace_user_id': repo.namespace_user.id,
-    'repository': repository,
-    'image_id': image_id,
-  }))
+  # Mark the image as uploaded.
+  _finish_image(namespace, repository, repo_image)
 
   return make_response('true', 200)