#434

backend/app/services/counts_export_download.rb

@@ -4,23 +4,70 @@ require 'fileutils'
 # Finish generation before reporting readiness; downloads never hold the full file in RAM.
 class CountsExportDownload
   EXPIRY = 10.minutes
+  MAX_CONCURRENT_EXPORTS = 2
+  MAX_FILE_BYTES = 128.megabytes
+  MAX_STORED_BYTES = 512.megabytes
+  MAX_EXPORT_SECONDS = 120
+  class Busy < StandardError; end
+  class TooLarge < StandardError; end
+  class TimedOut < StandardError; end
+  class StorageFull < StandardError; end
 
   def self.directory = Rails.root.join('tmp', 'counts_exports')
   def self.verifier = Rails.application.message_verifier('counts_export_download')
 
   def self.prepare(exporter, owner)
+    FileUtils.mkdir_p(directory, mode: 0o700)
+    slot = acquire_slot
+
+    begin
+      generate(exporter, owner)
+    ensure
+      slot.close
+    end
+  end
+
+  def self.acquire_slot
+    MAX_CONCURRENT_EXPORTS.times do |index|
+      path = directory.join(".slot-#{ index }.lock")
+      handle = File.open(path, File::RDWR | File::CREAT, 0o600)
+      return handle if handle.flock(File::LOCK_EX | File::LOCK_NB)
+
+      handle.close
+    end
+    raise Busy
+  end
+
+  def self.generate(exporter, owner)
     FileUtils.mkdir_p(directory, mode: 0o700)
     Dir.glob(directory.join('counts-*.export')).each do |path|
       File.unlink(path) if File.mtime(path) < Time.current - EXPIRY
     rescue Errno::ENOENT
       # Another request may have removed an expired download.
     end
+    stored = Dir.glob(directory.join('counts-*')).sum do |path|
+      File.size(path)
+    rescue Errno::ENOENT
+      0
+    end
+    reserved = MAX_CONCURRENT_EXPORTS * MAX_FILE_BYTES
+    raise StorageFull if stored + reserved > MAX_STORED_BYTES
+
     file = Tempfile.create(['counts-', '.part'], directory)
     complete = false
+    bytes = 0
+    deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + MAX_EXPORT_SECONDS
     path = file.path
     begin
       file.binmode
-      exporter.body.each { file.write(_1) }
+      exporter.body.each do |chunk|
+        raise TimedOut if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
+
+        bytes += chunk.bytesize
+        raise TooLarge if bytes > MAX_FILE_BYTES
+
+        file.write(chunk)
+      end
       file.close
       finished_path = path.sub(/\.part\z/, '.export')
       File.rename(path, finished_path)