Consolidate ChunkedIO

7 files changed, 122 insertions, 109 deletions

diff --git a/lib/gitlab/ci/trace/chunked_file/chunk_store/base.rb b/lib/gitlab/ci/trace/chunked_file/chunk_store/base.rb
index e2645918a40..6e104a6d764 100644
--- a/lib/gitlab/ci/trace/chunked_file/chunk_store/base.rb
+++ b/lib/gitlab/ci/trace/chunked_file/chunk_store/base.rb
@@ -22,18 +22,32 @@ module Gitlab
               raise NotImplementedError
             end
 
+            # Write data to chunk store. Always overwrite.
+            #
+            # @param [String] data
+            # @return [Fixnum] length of the data after writing
             def write!(data)
               raise NotImplementedError
             end
 
+            # Append data to chunk store
+            #
+            # @param [String] data
+            # @return [Fixnum] length of the appended
             def append!(data)
               raise NotImplementedError
             end
 
+            # Truncate data to chunk store
+            #
+            # @param [String] offset
             def truncate!(offset)
               raise NotImplementedError
             end
 
+            # Delete data from chunk store
+            #
+            # @param [String] offset
             def delete!
               raise NotImplementedError
             end
diff --git a/lib/gitlab/ci/trace/chunked_file/chunk_store/database.rb b/lib/gitlab/ci/trace/chunked_file/chunk_store/database.rb
index a7db214f428..3c2805e83f7 100644
--- a/lib/gitlab/ci/trace/chunked_file/chunk_store/database.rb
+++ b/lib/gitlab/ci/trace/chunked_file/chunk_store/database.rb
@@ -48,6 +48,8 @@ module Gitlab
             end
 
             def get
+              puts "#{self.class.name} - #{__callee__}: params[:chunk_index]: #{params[:chunk_index]}"
+
               job_trace_chunk.data
             end
 
@@ -56,9 +58,10 @@ module Gitlab
             end
 
             def write!(data)
+              raise NotImplementedError, 'Partial writing is not supported' unless params[:buffer_size] == data&.length
+              raise NotImplementedError, 'UPDATE (Overwriting data) is not supported' if job_trace_chunk.data
+
               puts "#{self.class.name} - #{__callee__}: data.length: #{data.length.inspect} params[:chunk_index]: #{params[:chunk_index]}"
-              raise NotImplementedError, 'Partial write is not supported' unless params[:buffer_size] == data&.length
-              raise NotImplementedError, 'UPDATE is not supported' if job_trace_chunk.data
 
               job_trace_chunk.data = data
               job_trace_chunk.save!
@@ -75,7 +78,10 @@ module Gitlab
             end
 
             def delete!
+              raise ActiveRecord::RecordNotFound, 'Could not find deletable record' unless job_trace_chunk.persisted?
+
               puts "#{self.class.name} - #{__callee__}: params[:chunk_index]: #{params[:chunk_index]}"
+
               job_trace_chunk.destroy!
             end
           end
diff --git a/lib/gitlab/ci/trace/chunked_file/chunk_store/redis.rb b/lib/gitlab/ci/trace/chunked_file/chunk_store/redis.rb
index cb45cd5fba5..c87275319d9 100644
--- a/lib/gitlab/ci/trace/chunked_file/chunk_store/redis.rb
+++ b/lib/gitlab/ci/trace/chunked_file/chunk_store/redis.rb
@@ -51,6 +51,9 @@ module Gitlab
               end
             end
 
+            BufferKeyNotFoundError = Class.new(StandardError)
+            WriteError = Class.new(StandardError)
+
             attr_reader :buffer_key
 
             def initialize(buffer_key, **params)
@@ -64,6 +67,8 @@ module Gitlab
             end
 
             def get
+              puts "#{self.class.name} - #{__callee__}: params[:chunk_index]: #{params[:chunk_index]}"
+
               Gitlab::Redis::Cache.with do |redis|
                 redis.get(buffer_key)
               end
@@ -76,35 +81,47 @@ module Gitlab
             end
 
             def write!(data)
+              raise ArgumentError, 'Could not write empty data' unless data.present?
+
               puts "#{self.class.name} - #{__callee__}: data.length: #{data.length.inspect} params[:chunk_index]: #{params[:chunk_index]}"
               Gitlab::Redis::Cache.with do |redis|
-                redis.set(buffer_key, data)
+                unless redis.set(buffer_key, data) == 'OK'
+                  raise WriteError, 'Failed to write'
+                end
+
                 redis.strlen(buffer_key)
               end
             end
 
             def append!(data)
+              raise ArgumentError, 'Could not write empty data' unless data.present?
+
               puts "#{self.class.name} - #{__callee__}: data.length: #{data.length.inspect} params[:chunk_index]: #{params[:chunk_index]}"
               Gitlab::Redis::Cache.with do |redis|
-                redis.append(buffer_key, data)
-                data.length
+                raise BufferKeyNotFoundError, 'Buffer key is not found' unless redis.exists(buffer_key)
+
+                original_size = size
+                new_size = redis.append(buffer_key, data)
+                appended_size = new_size - original_size
+
+                raise WriteError, 'Failed to append' unless appended_size == data.length
+
+                appended_size
               end
             end
 
             def truncate!(offset)
-              puts "#{self.class.name} - #{__callee__}: offset: #{offset.inspect} params[:chunk_index]: #{params[:chunk_index]}"
-              Gitlab::Redis::Cache.with do |redis|
-                return 0 unless redis.exists(buffer_key)
-
-                truncated_data = redis.getrange(buffer_key, 0, offset)
-                redis.set(buffer_key, truncated_data)
-              end
+              raise NotImplementedError
             end
 
             def delete!
               puts "#{self.class.name} - #{__callee__}: params[:chunk_index]: #{params[:chunk_index]}"
               Gitlab::Redis::Cache.with do |redis|
-                redis.del(buffer_key)
+                raise BufferKeyNotFoundError, 'Buffer key is not found' unless redis.exists(buffer_key)
+
+                unless redis.del(buffer_key) == 1
+                  raise WriteError, 'Failed to delete'
+                end
               end
             end
           end
diff --git a/lib/gitlab/ci/trace/chunked_file/chunked_io.rb b/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
index 5a2c56a687b..f28955264c8 100644
--- a/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
+++ b/lib/gitlab/ci/trace/chunked_file/chunked_io.rb
@@ -8,7 +8,7 @@ module Gitlab
     class Trace
       module ChunkedFile
         class ChunkedIO
-          extend ChunkedFile::Concerns::Opener
+          # extend ChunkedFile::Concerns::Opener
           include ChunkedFile::Concerns::Errors
           include ChunkedFile::Concerns::Hooks
           include ChunkedFile::Concerns::Callbacks
@@ -22,13 +22,21 @@ module Gitlab
 
           alias_method :pos, :tell
 
-          def initialize(job_id, size, mode = 'rb')
-            @size = size
+          def initialize(job_id, size = nil, mode = 'rb', &block)
+            raise NotImplementedError, "Mode 'w' is not supported" if mode.include?('w')
+
+            @size = size || calculate_size(job_id)
             @tell = 0
             @job_id = job_id
             @mode = mode
 
-            raise NotImplementedError, "Mode 'w' is not supported" if mode.include?('w')
+            if block_given?
+              begin
+                yield self
+              ensure
+                self.close
+              end
+            end
           end
 
           def close
@@ -128,7 +136,7 @@ module Gitlab
           end
 
           def present?
-            chunk_store.chunks_count(job_id) > 0
+            chunks_count > 0
           end
 
           def delete
@@ -177,19 +185,21 @@ module Gitlab
           end
 
           def write_chunk(data)
+            written_size = 0
+
             chunk_store.open(job_id, chunk_index, params_for_store) do |store|
               with_callbacks(:write_chunk, store) do
-                written_size = if buffer_size == data.length
+                written_size = if buffer_size == data.length || store.size == 0
                                  store.write!(data)
                                else
                                  store.append!(data)
                                end
 
                 raise WriteError, 'Written size mismatch' unless data.length == written_size
-
-                written_size
               end
             end
+
+            written_size
           end
 
           def truncate_chunk(offset)
@@ -228,7 +238,7 @@ module Gitlab
           end
 
           def chunks_count
-            (size / buffer_size)
+            (size / buffer_size.to_f).ceil
           end
 
           def first_chunk?
@@ -246,6 +256,10 @@ module Gitlab
           def buffer_size
             raise NotImplementedError
           end
+
+          def calculate_size(job_id)
+            chunk_store.chunks_size(job_id)
+          end
         end
       end
     end
diff --git a/lib/gitlab/ci/trace/chunked_file/concerns/hooks.rb b/lib/gitlab/ci/trace/chunked_file/concerns/hooks.rb
deleted file mode 100644
index 290a3a15805..00000000000
--- a/lib/gitlab/ci/trace/chunked_file/concerns/hooks.rb
+++ /dev/null
@@ -1,63 +0,0 @@
-module Gitlab
-  module Ci
-    class Trace
-      module ChunkedFile
-        module Concerns
-          module Hooks
-            extend ActiveSupport::Concern
-
-            included do
-              class_attribute :_before_methods, :_after_methods,
-                :instance_writer => false
-              self._before_methods = Hash.new []
-              self._after_methods = Hash.new []
-            end
-
-            class_methods do
-              def before_method(kind, callback)
-                self._before_methods = self._before_methods.
-                  merge kind => _before_methods[kind] + [callback]
-              end
-
-              def after_method(kind, callback)
-                self._after_methods = self._after_methods.
-                  merge kind => _after_methods[kind] + [callback]
-              end
-            end
-
-            def method_added(method_name)
-              return if self.class._before_methods.values.include?(method_name)
-              return if self.class._after_methods.values.include?(method_name)
-              return if hooked_methods.include?(method_name)
-
-              add_hooks_to(method_name)
-            end
-
-            private
-
-            def hooked_methods
-              @hooked_methods ||= []
-            end
-
-            def add_hooks_to(method_name)
-              hooked_methods << method_name
-
-              original_method = instance_method(method_name)
-
-              # re-define the method, but notice how we reference the original
-              # method definition
-              define_method(method_name) do |*args, &block|
-                self.class._before_methods[method_name].each { |hook| method(hook).call }
-
-                # now invoke the original method
-                original_method.bind(self).call(*args, &block).tap do
-                  self.class._after_methods[method_name].each { |hook| method(hook).call }
-                end
-              end
-            end
-          end
-        end
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/ci/trace/chunked_file/concerns/permissions.rb b/lib/gitlab/ci/trace/chunked_file/concerns/permissions.rb
index f364237f07b..11d8fd0cdfc 100644
--- a/lib/gitlab/ci/trace/chunked_file/concerns/permissions.rb
+++ b/lib/gitlab/ci/trace/chunked_file/concerns/permissions.rb
@@ -6,28 +6,19 @@ module Gitlab
           module Permissions
             extend ActiveSupport::Concern
 
+            WRITABLE_MODE = %w[a]
+            READABLE_MODE = %w[r +]
+
             included do
               PermissionError = Class.new(StandardError)
 
               attr_reader :write_lock_uuid
-
-              # mode checks
-              before_method :read, :can_read!
-              before_method :readline, :can_read!
-              before_method :each_line, :can_read!
-              before_method :write, :can_write!
-              before_method :truncate, :can_write!
-
-              # write_lock
-              before_method :write, :check_lock!
-              before_method :truncate, :check_lock!
-              before_method :delete, :check_lock!
             end
 
             def initialize(job_id, size, mode = 'rb')
-              if /(w|a)/ =~ mode
+              if WRITABLE_MODE.any? { |m| mode.include?(m) }
                 @write_lock_uuid = Gitlab::ExclusiveLease
-                  .new(write_lock_key, timeout: 1.hour.to_i).try_obtain
+                  .new(write_lock_key(job_id), timeout: 1.hour.to_i).try_obtain
 
                 raise PermissionError, 'Already opened by another process' unless write_lock_uuid
               end
@@ -37,25 +28,63 @@ module Gitlab
 
             def close
               if write_lock_uuid
-                Gitlab::ExclusiveLease.cancel(write_lock_key, write_lock_uuid)
+                Gitlab::ExclusiveLease.cancel(write_lock_key(job_id), write_lock_uuid)
               end
 
               super
             end
 
-            def check_lock!
-              raise PermissionError, 'Could not modify the file without lock' unless write_lock_uuid
+            def read(*args)
+              can_read!
+
+              super
+            end
+
+            def readline(*args)
+              can_read!
+
+              super
+            end
+
+            def each_line(*args)
+              can_read!
+
+              super
             end
 
+            def write(*args)
+              can_write!
+
+              super
+            end
+
+            def truncate(*args)
+              can_write!
+
+              super
+            end
+
+            def delete(*args)
+              can_write!
+
+              super
+            end
+
+            private
+
             def can_read!
-              raise IOError, 'not opened for reading' unless /(r|+)/ =~ mode
+              unless READABLE_MODE.any? { |m| mode.include?(m) }
+                raise IOError, 'not opened for reading'
+              end
             end
 
             def can_write!
-              raise IOError, 'not opened for writing' unless /(w|a)/ =~ mode
+              unless WRITABLE_MODE.any? { |m| mode.include?(m) }
+                raise IOError, 'not opened for writing'
+              end
             end
 
-            def write_lock_key
+            def write_lock_key(job_id)
               "live_trace:operation:write:#{job_id}"
             end
           end
diff --git a/lib/gitlab/ci/trace/chunked_file/live_trace.rb b/lib/gitlab/ci/trace/chunked_file/live_trace.rb
index 6976686553f..fe2205144cb 100644
--- a/lib/gitlab/ci/trace/chunked_file/live_trace.rb
+++ b/lib/gitlab/ci/trace/chunked_file/live_trace.rb
@@ -12,10 +12,6 @@ module Gitlab
 
           after_callback :write_chunk, :stash_to_database
 
-          def initialize(job_id, mode)
-            super(job_id, calculate_size(job_id), mode)
-          end
-
           def stash_to_database(store)
             # Once data is filled into redis, move the data to database
             if store.filled?