trd-rails-collector 0.1.0

data/README.rdoc

@@ -0,0 +1,35 @@
+= Treasure Data collector plugin for Rails
+
+= Getting Started
+
+Add the following line to your Gemfile:
+
+    gem 'trd-rails-collector'
+
+For Rails 2.x (not tested) without Bundler,
+edit +environment.rb+ and add to the initalizer block:
+
+    config.gem "trd-rails-collector"
+
+And then add +config/treasure_data.yml+ file as following:
+
+    defaults: &defaults
+      apikey: "YOUR_API_KEY"
+      database: myapp
+      table: access
+    
+    test:
+      <<: *defaults
+    
+    development:
+      <<: *defaults
+    
+    production:
+      <<: *defaults
+
+
+== Copyright
+
+Copyright:: Copyright (c) 2011 Treasure Data Inc.
+License::   Apache License, Version 2.0
+

data/lib/trd-rails-collector.rb

@@ -0,0 +1,298 @@
+
+module TreasureData
+
+
+  class AgentClass
+    def initialize
+      require 'thread'
+      require 'monitor'
+      require 'stringio'
+      require 'zlib'
+      require 'msgpack'
+      require 'time'
+      require 'net/http'
+      require 'cgi'
+      @map = {}  # (db,table) => buffer:String
+      @map.extend(MonitorMixin)
+      @cond = @map.new_cond
+      @queue = []
+      @finish = false
+      @time_keep_thread = nil
+      @upload_thread = nil
+    end
+
+    def init(apikey, database, table)
+      @apikey = apikey
+      @db = database
+      @table = table
+      unless @time_keep_thread
+        @time_keep_thread = Thread.new(&method(:time_keep_main))
+      end
+      unless @upload_thread
+        @upload_thread = Thread.new(&method(:upload_main))
+      end
+      nil
+    end
+
+    def init_rails(config)
+      require 'yaml'
+      sample = <<EOF
+  defaults: &defaults
+    apikey: "78fbbc1ff83088ab72e2b68c980c4c75786bee6e"
+    database: myapp
+    table: access
+
+  test:
+    <<: *defaults
+
+  development:
+    <<: *defaults
+
+  production:
+    <<: *defaults
+EOF
+      begin
+        yaml = YAML.load_file("#{RAILS_ROOT}/config/treasure_data.yml")
+      rescue
+        puts "Can't load config/treasure_data.yml file."
+        puts "  #{$!}"
+        puts "Put the following file:"
+        puts sample
+        return
+      end
+      env = yaml[RAILS_ENV]
+      unless env
+        puts "config/treasure_data.yml doesn't include setting for current environment (#{RAILS_ENV})."
+        return
+      end
+      apikey = env['apikey']
+      database = env['database']
+      table = env['table']
+      unless apikey && database && table
+        puts "config/treasure_data.yml doesn't include setting for current environment (#{RAILS_ENV})."
+        return
+      end
+      init(apikey, database, table)
+      config.middleware.use TreasureData::Rack
+    end
+
+    SIZE_LIMIT = 1024*1024  # 1MB
+    TIME_LIMIT = 5         # 1min TODO
+    SLEEP_TIME = 10         # 10sec TODO
+
+    API_HOST = 'api.treasure-data.com'
+    API_PORT = 80
+    USE_SSL = false
+
+    def emit(record)
+      now = Time.now.to_i
+      record['time'] ||= now
+
+      key = [@db, @table]
+
+      @map.synchronize do
+
+        buffer = @map[key]
+        unless buffer
+          buffer = ''
+          time = Time.now.to_i
+          (class<<buffer;self;end).module_eval do
+            define_method(:creation_time) do
+              time
+            end
+          end
+          @map[key] = buffer
+        end
+
+        buffer = (@map[key] ||= '')
+        record.to_msgpack(buffer)
+
+        if buffer.size > SIZE_LIMIT
+          @queue << [@db, @table, buffer]
+          @map.delete(key)
+          @cond.signal
+        end
+
+      end
+      nil
+    end
+
+    private
+    def time_keep_main
+      until @finish
+        sleep SLEEP_TIME
+        now = Time.now.to_i
+
+        @map.synchronize do
+          one = false
+          @map.delete_if {|db_table,buffer|
+            if now - buffer.creation_time > TIME_LIMIT
+              @queue << [db_table[0], db_table[1], buffer]
+              one = true
+            end
+          }
+          @cond.broadcast if one
+        end
+      end
+    end
+
+    def upload_main
+      tuple = nil
+
+      until @finish
+        @map.synchronize do
+          if tuple
+            @queue.shift
+            tuple = nil
+          end
+
+          while true
+            break if @finish
+            unless @queue.empty?
+              tuple = @queue[0]
+              break
+            end
+            @cond.wait
+          end
+        end
+
+        break if @finish
+
+        begin
+          code, body = upload(*tuple)
+        rescue
+          # TODO retry
+          $stderr.puts $!
+        end
+      end
+
+      # TODO upload_all
+    end
+
+    def upload(db, table, buffer)
+      str = StringIO.new
+      Zlib::GzipWriter.wrap(str) {|gz|
+        gz.write buffer
+      }
+      data = str.string
+
+      http = Net::HTTP.new(API_HOST, API_PORT)
+      if USE_SSL
+        http.use_ssl = true
+        http.verify_mode = OpenSSL::SSL::VERIFY_PEER
+        store = OpenSSL::X509::Store.new
+        http.cert_store = store
+      end
+
+      header = {}
+      header['Authorization'] = "TRD #{@apikey}"
+      header['Date'] = Time.now.rfc2822
+      header['Content-Length'] = data.size.to_s
+
+      url = "/v2/import/#{e @db}/#{e @table}/msgpack.gz"
+
+      RAILS_DEFAULT_LOGGER.debug "Uploading compressed logs #{data.size} bytes to #{url}"
+      request = Net::HTTP::Put.new(url, header)
+      request.body = data
+
+      response = http.request(request)
+      return response.code, response.body
+    end
+
+    def e(s)
+      CGI.escape(s.to_s)
+    end
+  end
+
+  Agent = AgentClass.new
+
+
+  def self.init(apikey, database, table)
+    Agent.init(apikey, database, table)
+  end
+
+  def self.init_rails(config)
+    Agent.init_rails(config)
+  end
+
+  def self.emit(record)
+    Agent.emit(record)
+  end
+
+
+  class Rack
+    PARAM_KEYS = [
+      # Rails 3
+      'action_dispatch.request.path_parameters',
+
+      # Rack default
+      'rack.routing_args',
+    ]
+
+    def initialize(app, options={})
+      @app = app
+    end
+
+    def call(env)
+      r = @app.call(env)
+
+      m = {}
+
+      # compatibility for fluent's apache log parser
+      m['host'] = env['REMOTE_HOST']
+      m['method'] = env['REQUEST_METHOD']
+      # TODO m['user']
+      m['path'] = env['PATH_INFO']
+      m['code'] = r[0].to_i
+      m['size'] = r[1]['Content-Length'].to_i
+      m['referer'] = env['HTTP_REFERER'] || '-'
+      m['agent'] = env['HTTP_USER_AGENT']
+
+      # additional information
+      m['server'] = env['SERVER_NAME']
+
+      # parameters
+      PARAM_KEYS.each {|key|
+        if e = env[key]
+          # TODO namespace: add param_ prefix?
+          m.merge!(e)
+        end
+      }
+
+      # time parameter is required
+      m['time'] = Time.now.to_i
+
+      TreasureData.emit(m)
+
+      #File.open('out.log', "a") {|f|
+      #  f.write m.to_json+"\n"
+      #}
+
+      #File.open('envlog.txt', "a") {|f|
+      #  f.write env.pretty_inspect
+      #}
+
+      r
+    end
+  end
+end
+
+
+if defined? Rails
+  if Rails.respond_to?(:version) && Rails.version =~ /^3/
+    module TreasureData
+      class Railtie < Rails::Railtie
+        initializer "treasure_data_agent.start_plugin" do |app|
+          TreasureData.init_rails(app.config)
+        end
+      end
+    end
+  else
+    # After verison 2.0 of Rails we can access the configuration directly.
+    #if Rails.respond_to?(:configuration)
+    #  not available for Rails versions prior to 2.2
+    #end
+    TreasureData.init_rails(Rails.configuration)
+  end
+end
+

metadata

@@ -0,0 +1,83 @@
+--- !ruby/object:Gem::Specification 
+name: trd-rails-collector
+version: !ruby/object:Gem::Version 
+  hash: 27
+  prerelease: false
+  segments: 
+  - 0
+  - 1
+  - 0
+  version: 0.1.0
+platform: ruby
+authors: 
+- Sadayuki Furuhashi
+autorequire: 
+bindir: bin
+cert_chain: []
+
+date: 2011-06-25 00:00:00 +09:00
+default_executable: 
+dependencies: 
+- !ruby/object:Gem::Dependency 
+  name: msgpack
+  prerelease: false
+  requirement: &id001 !ruby/object:Gem::Requirement 
+    none: false
+    requirements: 
+    - - ~>
+      - !ruby/object:Gem::Version 
+        hash: 7
+        segments: 
+        - 0
+        - 4
+        - 4
+        version: 0.4.4
+  type: :runtime
+  version_requirements: *id001
+description: 
+email: 
+executables: []
+
+extensions: []
+
+extra_rdoc_files: 
+- README.rdoc
+files: 
+- lib/trd-rails-collector.rb
+- README.rdoc
+has_rdoc: true
+homepage: 
+licenses: []
+
+post_install_message: 
+rdoc_options: 
+- --charset=UTF-8
+require_paths: 
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 3
+      segments: 
+      - 0
+      version: "0"
+required_rubygems_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 3
+      segments: 
+      - 0
+      version: "0"
+requirements: []
+
+rubyforge_project: 
+rubygems_version: 1.3.7
+signing_key: 
+specification_version: 3
+summary: Treasure Data collector plugin for Rails
+test_files: []
+