From d3c135ec7261f73a8ba9c67cf6ac65bb472cd5b6 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Tue, 18 Sep 2018 02:10:20 -0500
Subject: [PATCH] Downloads::File#http_get_streaming: clean up retry logic.

Replace handrolled retry logic with retriable gem (already pulled in by another gem).
---
 Gemfile                          |  1 +
 Gemfile.lock                     |  1 +
 app/logical/downloads/file.rb    | 59 +++++++++++++-------------------
 test/unit/downloads/file_test.rb | 11 ++----
 4 files changed, 29 insertions(+), 43 deletions(-)

diff --git a/Gemfile b/Gemfile
index c79a1144b..717058a88 100644
--- a/Gemfile
+++ b/Gemfile
@@ -50,6 +50,7 @@ gem 'activemodel-serializers-xml'
 gem 'ptools'
 gem 'jquery-rails'
 gem 'webpacker', '>= 4.0.x'
+gem 'retriable'
 
 # needed for looser jpeg header compat
 gem 'ruby-imagespec', :require => "image_spec", :git => "https://github.com/r888888888/ruby-imagespec.git", :branch => "exif-fixes"
diff --git a/Gemfile.lock b/Gemfile.lock
index 0a9a2b049..3b75b6a31 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -489,6 +489,7 @@ DEPENDENCIES
   rakismet
   recaptcha
   responders
+  retriable
   ruby-imagespec!
   ruby-prof
   ruby-vips
diff --git a/app/logical/downloads/file.rb b/app/logical/downloads/file.rb
index c8bbc0b78..eb2aece9c 100644
--- a/app/logical/downloads/file.rb
+++ b/app/logical/downloads/file.rb
@@ -2,6 +2,8 @@ module Downloads
   class File
     class Error < Exception ; end
 
+    RETRIABLE_ERRORS = [Errno::ECONNRESET, Errno::ETIMEDOUT, Errno::EIO, Errno::EHOSTUNREACH, Errno::ECONNREFUSED, Timeout::Error, IOError]
+
     attr_reader :data, :options
     attr_accessor :source, :referer
 
@@ -52,16 +54,15 @@ module Downloads
       end
     end
 
-    def download!
+    def download!(tries: 3)
       strategy = Sources::Strategies.find(source, referer)
       output_file = Tempfile.new(binmode: true)
+      url = self.class.uncached_url(strategy.file_url, strategy.headers)
       @data = strategy.data
 
-      http_get_streaming(
-        self.class.uncached_url(strategy.file_url, strategy.headers), 
-        output_file, 
-        strategy.headers
-      )
+      Retriable.retriable(on: RETRIABLE_ERRORS, tries: tries, base_interval: 0) do
+        http_get_streaming(url, output_file, strategy.headers)
+      end
 
       [output_file, strategy]
     end
@@ -74,42 +75,30 @@ module Downloads
     end
 
     def http_get_streaming(src, file, headers = {}, max_size: Danbooru.config.max_file_size)
-      tries = 0
       url = URI.parse(src)
 
-      while true
-        unless url.is_a?(URI::HTTP) || url.is_a?(URI::HTTPS)
-          raise Error.new("URL must be HTTP or HTTPS")
-        end
+      unless url.is_a?(URI::HTTP) || url.is_a?(URI::HTTPS)
+        raise Error.new("URL must be HTTP or HTTPS")
+      end
 
-        validate_local_hosts(url)
+      validate_local_hosts(url)
 
-        begin
-          size = 0
-          options = { stream_body: true, timeout: 10, headers: headers }
+      size = 0
+      options = { stream_body: true, timeout: 10, headers: headers }
 
-          res = HTTParty.get(url, options.deep_merge(Danbooru.config.httparty_options)) do |chunk|
-            size += chunk.size
-            raise Error.new("File is too large (max size: #{max_size})") if size > max_size && max_size > 0
+      res = HTTParty.get(url, options.deep_merge(Danbooru.config.httparty_options)) do |chunk|
+        size += chunk.size
+        raise Error.new("File is too large (max size: #{max_size})") if size > max_size && max_size > 0
 
-            file.write(chunk)
-          end
+        file.write(chunk)
+      end
 
-          if res.success?
-            file.rewind
-            return file
-          else
-            raise Error.new("HTTP error code: #{res.code} #{res.message}")
-          end
-        rescue Errno::ECONNRESET, Errno::ETIMEDOUT, Errno::EIO, Errno::EHOSTUNREACH, Errno::ECONNREFUSED, Timeout::Error, IOError => x
-          tries += 1
-          if tries < 3
-            retry
-          else
-            raise
-          end
-        end
-      end # while
+      if res.success?
+        file.rewind
+        return file
+      else
+        raise Error.new("HTTP error code: #{res.code} #{res.message}")
+      end
     end # def
   end
 end
diff --git a/test/unit/downloads/file_test.rb b/test/unit/downloads/file_test.rb
index 9ce6fc808..435acacd9 100644
--- a/test/unit/downloads/file_test.rb
+++ b/test/unit/downloads/file_test.rb
@@ -10,14 +10,9 @@ module Downloads
       end
 
       context "that fails" do
-        setup do
-          HTTParty.stubs(:get).raises(Errno::ETIMEDOUT)
-        end
-
-        should "retry three times" do
-          assert_raises(Errno::ETIMEDOUT) do
-            @download.http_get_streaming(@source, @tempfile)
-          end
+        should "retry three times before giving up" do
+          HTTParty.expects(:get).times(3).raises(Errno::ETIMEDOUT)
+          assert_raises(Errno::ETIMEDOUT) { @download.download! }
         end
       end