S3 Support for files

.rbenv-version

@@ -1 +1,2 @@
+1.9.3-p194
 jruby-1.7.0-dev

Gemfile

@@ -1,7 +1,9 @@
-source :rubygems
+source 'https://rubygems.org'
 
 gem 'eventmachine', '~>1.0.0.rc.4'
 gem 'redis', "~> 3.0.2"
+gem 'aws-s3', "~> 0.6.3"
+
 if RUBY_PLATFORM == 'java'
   gem 'json-jruby'
 else

Gemfile.lock

@@ -1,26 +1,32 @@
 GEM
-  remote: http://rubygems.org/
+  remote: https://rubygems.org/
   specs:
+    aws-s3 (0.6.3)
+      builder
+      mime-types
+      xml-simple
+    builder (3.0.4)
     eventmachine (1.0.0.rc.4)
     eventmachine (1.0.0.rc.4-java)
     json (1.5.0)
     json (1.5.0-java)
-    json-jruby (1.5.0-java)
-      json (= 1.5.0)
     metaclass (0.0.1)
+    mime-types (1.23)
     mocha (0.11.3)
       metaclass (~> 0.0.1)
     rake (0.9.2.2)
     redis (3.0.2)
+    xml-simple (1.1.2)
     yard (0.8.2)
 
 PLATFORMS
   java
   ruby
 
 DEPENDENCIES
+  aws-s3 (~> 0.6.3)
   eventmachine (~> 1.0.0.rc.4)
-  json-jruby
+  json
   mocha
   rake
   redis (~> 3.0.2)

README.md

@@ -7,13 +7,10 @@ implementation](https://github.com/etsy/statsd), which they described in
 a [blog post](http://codeascraft.etsy.com/2011/02/15/measure-anything-measure-everything/).
 
 Batsd differs from etsy's statsd implementation primarily in how it stores data
--- data is stored to a combination of Redis and flat files on disk. You can
-read more about persistence in [About:
-Persistence](http://noahhl.github.com/batsd/doc/file.persistence.html).
+-- data is stored to a combination of Redis and flat files on disk (or AWS S3). You can
+read more about persistence in [About: Persistence](http://noahhl.github.com/batsd/doc/file.persistence.html).
 
-Batsd grew out of usage at [37signals](http://37signals.com), where it has been
-used for the last year. An [earlier form](https://github.com/noahhl/statsd-server) was
-inspired by [quasor](https://github.com/quasor/statsd). 
+Batsd grew out of usage at [37signals](http://37signals.com), where it has been used for the last year. An [earlier form](https://github.com/noahhl/statsd-server) was inspired by [quasor](https://github.com/quasor/statsd). 
 
 ### Documentation:
 
@@ -50,11 +47,26 @@ Example config.yml
     # Where to store data. Data at the first retention level is stored
     # in redis; further data retentions are stored on disk
 
-    # Root path to store disk aggregations
-    root: /statsd 
     redis:
       :host: 127.0.0.1
       :port: 6379
+
+    # Which file store to use. Available are: diskstore and s3
+    filestore: 'diskstore'
+
+
+    # Diskstore configuration. 
+    # This is only if you have filestore set to diskstore
+    diskstore:
+      # Root path to store disk aggregations
+      :root: tmp/statsd 
+
+    # S3 authentication and storage information
+    # This is only if you have filestore set to s3
+    s3:
+      :access_key: 'access key goes here'
+      :secret_access_key: 'secret key goes here'
+      :bucket: 'bucket name goes here'
 
     # Configure how much data to retain at what intervals
     # Key is seconds, value is number of measurements at that

config.yml

@@ -1,11 +1,23 @@
 bind: 0.0.0.0
 port: 8125
-root: tmp/statsd
+
+filestore: 'diskstore'
+
+s3:
+  :access_key: 'access key goes here'
+  :secret_access_key: 'secret key goes here'
+  :bucket: 'bucket name goes here'
+
+diskstore:
+  :root: tmp/statsd
+
 redis:
   host: 127.0.0.1
   port: 6379
+
 retentions:
   10:  360 #1 hour
   60:  10080  #1 week
   600: 52594 #1 year
+
 autotruncate: false

lib/batsd.rb

@@ -1,10 +1,13 @@
 require 'benchmark'
 require 'eventmachine'
 require 'redis'
+require 'aws/s3'
 
 require 'core-ext/array'
 
-require 'batsd/diskstore'
+require 'batsd/filestore'
+require 'batsd/filestore/diskstore'
+require 'batsd/filestore/s3'
 require 'batsd/redis'
 require 'batsd/threadpool'
 require 'batsd/receiver'

lib/batsd/deleter.rb

@@ -6,13 +6,13 @@ class Deleter
 
     # Create a new deleter
     #
-    # * Establish the diskstore that will be used
+    # * Establish the filestore that will be used
     # * Establish the redis connection that will be needed
     #
     def initialize(options={})
       @options = options
-      @redis = Batsd::Redis.new(options )
-      @diskstore = Batsd::Diskstore.new(options[:root])
+      @redis = Batsd::Redis.new(options)
+      @filestore = Batsd::Filestore.init(options)
     end
 
     def delete(statistic)
@@ -37,7 +37,7 @@ def delete(statistic)
         # other retentions
         retentions.each do |retention|
           key = "#{statistic}:#{retention}"
-          @diskstore.delete(@diskstore.build_filename(key), :delete_empty_dirs => true)
+          @filestore.delete(@filestore.build_filename(key), delete_empty_dirs: true)
         end
       end
       deletions

lib/batsd/filestore.rb

@@ -0,0 +1,35 @@
+module Batsd
+
+  # A wrapper around all file stores (Diskstore and S3), includes commonality and control functions
+  class Filestore
+    attr_accessor :root
+
+    # Creates an instance of the right child depending on configuration
+    def self.init(options)
+
+      unless options[:filestore]
+        Batsd::Diskstore.new diskstore: { root: (options[:root] || options[:diskstore][:root]) }
+      else
+        case options[:filestore].downcase
+        when 'diskstore' then
+          Batsd::Diskstore.new(options)
+        when 's3' then
+          Batsd::S3.new(options)
+        end
+      end
+    end
+
+    def build_filename(statistic)
+      return unless statistic
+      paths     = []
+      file_hash = Digest::MD5.hexdigest(statistic)
+
+      paths << @root if @root
+      paths << file_hash[0,2]
+      paths << file_hash[2,2]
+      paths << file_hash
+
+      File.join(paths)
+    end
+  end
+end

lib/batsd/diskstore.rb → lib/batsd/filestore/diskstore.rb

@@ -1,26 +1,11 @@
-require 'digest'
 require 'fileutils'
-
 module Batsd 
   # Handles disk operations -- writing, truncating, and reading
-  class Diskstore
+  class Diskstore < Filestore
 
     # Create a new diskstore object
-    def initialize(root)
-      @root = root
-    end
-
-    # Calculate the filename that will be used to store the
-    # metric to disk.
-    #
-    # Filenames are MD5 hashes of the statistic name, including any
-    # aggregation-based suffix, and are stored in two levels of nested
-    # directories (e.g., <code>/00/01/0001s0d03dd0s030d03d</code>)
-    #
-    def build_filename(statistic)
-      return unless statistic
-      file_hash = Digest::MD5.hexdigest(statistic)
-      File.join(@root, file_hash[0,2], file_hash[2,2], file_hash)
+    def initialize(options)
+      @root = options[:diskstore][:root]
     end
 
     # Append a value to a file

lib/batsd/filestore/s3.rb

@@ -0,0 +1,110 @@
+module Batsd
+  # Handles writing, truncating, and reading on AWS S3 buckets
+  # Meant to be a replacement for Batsd::Diskstore for usage on heroku
+  class S3 < Filestore
+    S3_FOLDER_EXT = '_$folder$'
+
+    def initialize(options)
+      @credentials = options[:s3]
+      @bucket      = options[:s3][:bucket]
+    end
+
+    # Fetch a file from AWS S3
+    def fetch_file(filename)
+      establish_connection
+      data = AWS::S3::S3Object.value(filename, @bucket) if AWS::S3::S3Object.exists?(filename, @bucket) 
+
+      data || ''
+    end
+
+
+    # Store a file to AWS S3
+    # folders and file types are handled by the gem
+
+    def store_file(filename, file_data)
+      establish_connection
+
+      AWS::S3::S3Object.store filename, StringIO.new(file_data), @bucket, access: :authenticated_read
+    end
+
+    # Append a value to a file
+    #
+    # Open the file in append mode (creating directories needed along
+    # the way), write the value and a newline, and close the file again.
+    #
+    def append_value_to_file(filename, value, attempts=0)
+      file_data  = fetch_file(filename) + "#{value}\n"
+
+      store_file filename, file_data
+    rescue Exception => e
+      puts "Encountered an error trying to store to #{filename}: #{e} #{e.message} #{e.backtrace if ENV["VERBOSE"]}"
+      if attempts < 2
+        puts "Retrying #{filename} for the #{attempts+1} time"
+        append_value_to_file(filename, value, attempts+1)
+      end
+    end
+
+    # Reads the set of values in the range desired from file
+    #
+    # Reads until it reaches end_ts or the end fo the file. Returns an array
+    # of <code>{timestamp: ts, value: v}</code> hashes.
+    #
+    def read(statistic, start_ts, end_ts)
+      datapoints = []
+      filename = build_filename statistic
+
+      begin
+        file_data = fetch_file(filename)
+        file_data.split("\n").each do |line|
+          ts, value = line.split
+          if ts >= start_ts && ts <= end_ts
+            datapoints << { timestamp: ts.to_i, value: value }
+          end
+        end
+      rescue Exception => e
+        puts "Encountered an error trying to read #{filename}: #{e} #{e.message} #{e.backtrace if ENV["VERBOSE"]}"
+      end
+
+      datapoints
+    end
+
+    # Truncates a file by rewriting to a temp file everything after the since
+    # timestamp that is provided. The temp file is then renaemed to the
+    # original.
+    #
+    def truncate(filename, since)
+      puts "Truncating #{filename} since #{since}" if ENV["VVERBOSE"]
+
+      truncated_file_data = ''
+      old_file_data       = fetch_file(filename)
+
+      old_file_data.split("\n").each do |line|
+        truncated_file_data += "#{line}\n" if(line.split[0] >= since rescue true)
+      end
+
+      store_file filename, truncated_file_data
+
+      truncated_file_data
+    rescue Exception => e
+      puts "Encountered an error trying to truncate #{filename}: #{e} #{e.message} #{e.backtrace if ENV["VERBOSE"]}"
+    end
+
+    # Deletes a file, if it exists.
+    # If :delete_empty_dirs is true, empty directories will be deleted too.
+    # TODO: Support for :delete_empty_dirs
+    #
+    def delete(filename, options={})
+      establish_connection
+
+      AWS::S3::S3Object.find(filename, @bucket).delete if AWS::S3::S3Object.exists? filename, @bucket
+    rescue Exception => e
+      puts "Encountered an error trying to delete #{filename}: #{e} #{e.message} #{e.backtrace if ENV["VERBOSE"]}"
+    end
+
+    def establish_connection
+      unless AWS::S3::Base.connected?
+        AWS::S3::Base.establish_connection! access_key_id: @credentials[:access_key], secret_access_key: @credentials[:secret_access_key]
+      end
+    end
+  end
+end

lib/batsd/handlers/counter.rb

@@ -10,17 +10,18 @@ class Handler::Counter < Handler
     # Set up a new handler to handle counters
     #
     # * Set up a redis client
-    # * Set up a diskstore client to write aggregates to disk
+    # * Set up a filestore client to write aggregates to disk
     # * Initialize last flush timers to now
     #
     def initialize(options)
-      @redis = Batsd::Redis.new(options)
-      @diskstore = Batsd::Diskstore.new(options[:root])
-      @counters = @active_counters = {}
-      @retentions = options[:retentions].keys
-      @flush_interval = @retentions.first
       now = Time.now.to_i
-      @last_flushes = @retentions.inject({}){|l, r| l[r] = now; l }
+
+      @redis          = Batsd::Redis.new(options)
+      @filestore      = Batsd::Filestore.init(options)
+      @counters       = @active_counters = {}
+      @retentions     = options[:retentions].keys
+      @flush_interval = @retentions.first
+      @last_flushes   = @retentions.inject({}){|l, r| l[r] = now; l }
       super
     end
 
@@ -101,7 +102,7 @@ def flush
                   value = @redis.get_and_clear_key(key)
                   if value
                     value = "#{ts} #{value}"
-                    @diskstore.append_value_to_file(@diskstore.build_filename(key), value)
+                    @filestore.append_value_to_file(@filestore.build_filename(key), value)
                   end
                 end
               end

lib/batsd/handlers/gauge.rb

@@ -12,11 +12,11 @@ class Handler::Gauge < Handler
     # Set up a new handler to handle gauges
     #
     # * Set up a redis client
-    # * Set up a diskstore client to write aggregates to disk
+    # * Set up a filestore client to write aggregates to disk
     #
     def initialize(options)
-      @redis = Batsd::Redis.new(options)
-      @diskstore = Batsd::Diskstore.new(options[:root])
+      @redis     = Batsd::Redis.new(options)
+      @filestore = Batsd::Filestore.init(options)
       super
     end
 
@@ -34,7 +34,7 @@ def handle(key, value, sample_rate)
         end
         value = "#{timestamp} #{value}"
         key   = "gauges:#{key}"
-        @diskstore.append_value_to_file(@diskstore.build_filename(key), value)
+        @filestore.append_value_to_file(@filestore.build_filename(key), value)
         @redis.add_datapoint key
       end
     end