Add latest changes from gitlab-org/gitlab@16-1-stable-eev16.1.0-rc42

2 files changed, 128 insertions, 11 deletions

diff --git a/lib/object_storage/fog_helpers.rb b/lib/object_storage/fog_helpers.rb
new file mode 100644
index 00000000000..1db75ea24b9
--- /dev/null
+++ b/lib/object_storage/fog_helpers.rb
@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module ObjectStorage
+  module FogHelpers
+    include ::Gitlab::Utils::StrongMemoize
+
+    def available?
+      object_store.enabled
+    end
+
+    private
+
+    def delete_object(key)
+      return unless available?
+
+      connection.delete_object(bucket_name, object_key(key))
+
+    # So far, only GoogleCloudStorage raises an exception when the file is not found.
+    # Other providers support idempotent requests and does not raise an error
+    # when the file is missing.
+    rescue ::Google::Apis::ClientError => e
+      Gitlab::ErrorTracking.log_exception(e)
+    end
+
+    def storage_location_identifier
+      raise NotImplementedError, "#{self} does not implement #{__method__}"
+    end
+
+    def object_store
+      ObjectStorage::Config::LOCATIONS.fetch(storage_location_identifier).object_store
+    end
+
+    def bucket_name
+      object_store.remote_directory
+    end
+
+    def object_key(key)
+      # We allow administrators to create "sub buckets" by setting a prefix.
+      # This makes it possible to deploy GitLab with only one object storage
+      # bucket. This mirrors the implementation in app/uploaders/object_storage.rb.
+      File.join([object_store.bucket_prefix, key].compact)
+    end
+
+    def connection
+      return unless available?
+
+      ::Fog::Storage.new(object_store.connection.to_hash.deep_symbolize_keys)
+    end
+    strong_memoize_attr :connection
+  end
+end
diff --git a/lib/object_storage/pending_direct_upload.rb b/lib/object_storage/pending_direct_upload.rb
index 3e84bc4ebc9..3a930e0e0af 100644
--- a/lib/object_storage/pending_direct_upload.rb
+++ b/lib/object_storage/pending_direct_upload.rb
@@ -2,31 +2,97 @@
 
 module ObjectStorage
   class PendingDirectUpload
+    include ObjectStorage::FogHelpers
+
     KEY = 'pending_direct_uploads'
+    MAX_UPLOAD_DURATION = 3.hours.freeze
 
-    def self.prepare(location_identifier, path)
-      ::Gitlab::Redis::SharedState.with do |redis|
+    def self.prepare(location_identifier, object_storage_path)
+      with_redis do |redis|
         # We need to store the location_identifier together with the timestamp to properly delete
         # this object if ever this upload gets stale. The location identifier will be used
         # by the clean up worker to properly generate the storage options through ObjectStorage::Config.for_location
-        redis.hset(KEY, key(location_identifier, path), Time.current.utc.to_i)
+        key = redis_key(location_identifier, object_storage_path)
+        redis.hset(KEY, key, Time.current.utc.to_i)
+        log_event(:prepared, key)
+      end
+    end
+
+    def self.exists?(location_identifier, object_storage_path)
+      with_redis do |redis|
+        redis.hexists(KEY, redis_key(location_identifier, object_storage_path))
+      end
+    end
+
+    def self.complete(location_identifier, object_storage_path)
+      with_redis do |redis|
+        key = redis_key(location_identifier, object_storage_path)
+        redis.hdel(KEY, key)
+        log_event(:completed, key)
       end
     end
 
-    def self.exists?(location_identifier, path)
-      ::Gitlab::Redis::SharedState.with do |redis|
-        redis.hexists(KEY, key(location_identifier, path))
+    def self.redis_key(location_identifier, object_storage_path)
+      [location_identifier, object_storage_path].join(':')
+    end
+
+    def self.count
+      with_redis do |redis|
+        redis.hlen(KEY)
       end
     end
 
-    def self.complete(location_identifier, path)
-      ::Gitlab::Redis::SharedState.with do |redis|
-        redis.hdel(KEY, key(location_identifier, path))
+    def self.each
+      with_redis do |redis|
+        redis.hscan_each(KEY) do |entry|
+          redis_key, timestamp = entry
+          storage_location_identifier, object_storage_path = redis_key.split(':')
+
+          object = new(
+            redis_key: redis_key,
+            storage_location_identifier: storage_location_identifier,
+            object_storage_path: object_storage_path,
+            timestamp: timestamp
+          )
+
+          yield(object)
+        end
       end
     end
 
-    def self.key(location_identifier, path)
-      [location_identifier, path].join(':')
+    def self.with_redis(&block)
+      Gitlab::Redis::SharedState.with(&block) # rubocop:disable CodeReuse/ActiveRecord
     end
+
+    def self.log_event(event, redis_key)
+      Gitlab::AppLogger.info(
+        message: "Pending direct upload #{event}",
+        redis_key: redis_key
+      )
+    end
+
+    def initialize(redis_key:, storage_location_identifier:, object_storage_path:, timestamp:)
+      @redis_key = redis_key
+      @storage_location_identifier = storage_location_identifier.to_sym
+      @object_storage_path = object_storage_path
+      @timestamp = timestamp.to_i
+    end
+
+    def stale?
+      timestamp < MAX_UPLOAD_DURATION.ago.utc.to_i
+    end
+
+    def delete
+      delete_object(object_storage_path)
+
+      self.class.with_redis do |redis|
+        redis.hdel(self.class::KEY, redis_key)
+        self.class.log_event(:deleted, redis_key)
+      end
+    end
+
+    private
+
+    attr_reader :redis_key, :storage_location_identifier, :object_storage_path, :timestamp
   end
 end