distribot 0.1.1

data/lib/distribot.rb

@@ -0,0 +1,108 @@
+
+require 'active_support/core_ext/object'
+require 'active_support/json'
+require 'redis'
+require 'distribot/flow'
+require 'distribot/phase'
+require 'distribot/phase_handler'
+require 'distribot/handler'
+require 'distribot/flow_created_handler'
+require 'distribot/phase_started_handler'
+require 'distribot/worker'
+require 'distribot/task_finished_handler'
+require 'distribot/handler_finished_handler'
+require 'distribot/phase_finished_handler'
+require 'distribot/flow_finished_handler'
+require 'distribot/connector'
+require 'syslog/logger'
+require 'logstash-logger'
+require 'concurrent'
+
+module Distribot
+  class << self
+    attr_accessor :config, :did_configure, :connector, :redis, :debug, :logger
+    @config = OpenStruct.new
+    @did_configure = false
+    @connector = nil
+
+    def reset_configuration!
+      self.config = OpenStruct.new
+      self.did_configure = false
+      self.redis = nil
+    end
+
+    def configure(&block)
+      reset_configuration!
+      @did_configure = true
+      block.call(configuration)
+      # Now set defaults for things that aren't defined:
+      configuration.redis_prefix ||= 'distribot'
+      configuration.queue_prefix ||= 'distribot'
+    end
+
+    def connector
+      @connector ||= BunnyConnector.new(configuration.rabbitmq_url)
+    end
+
+    def configuration
+      unless @did_configure
+        reset_configuration!
+        configure do |config|
+          config.redis_url = ENV['DISTRIBOT_REDIS_URL']
+          config.rabbitmq_url = ENV['DISTRIBOT_RABBITMQ_URL']
+        end
+      end
+      self.config
+    end
+
+    def queue_exists?(name)
+      connector.queue_exists?(name)
+    end
+
+    def redis
+      # Redis complains if we pass it a nil url. Better to not pass a url:
+      @redis ||= if configuration.redis_url
+                   Redis.new(url: configuration.redis_url)
+                 else
+                   Redis.new
+                 end
+    end
+
+    def debug=(value)
+      @debug = value ? true : false
+    end
+
+    def debug
+      @debug ||= false
+    end
+
+    def redis_id(type, id)
+      "#{configuration.redis_prefix}-#{type}:#{id}"
+    end
+
+    def publish!(topic, data)
+      connector.publish(topic, data)
+    end
+
+    def subscribe(topic, options = {}, &block)
+      connector.subscribe(topic, options) do |message|
+        block.call(message)
+      end
+    end
+
+    def broadcast!(topic, data)
+      connector.broadcast(topic, data)
+    end
+
+    def subscribe_multi(topic, options = {}, &block)
+      connector.subscribe_multi(topic, options) do |message|
+        block.call(message)
+      end
+    end
+
+    def logger
+      @logger ||= LogStashLogger.new(type: :syslog, formatter: :json)
+      @logger
+    end
+  end
+end

data/provision/nodes.sh

@@ -0,0 +1,80 @@
+#!/bin/bash
+
+set -e
+
+setup_dependencies() {
+  sudo apt-get -y update
+  sudo apt-get -y autoremove
+  sudo apt-get install -y \
+    ruby2.0 \
+    ruby2.0-dev \
+    build-essential \
+    git \
+    wget \
+    vim \
+    librabbitmq-dev \
+    psmisc \
+    curl \
+    libcurl4-gnutls-dev \
+    python
+
+  sudo ln -sf /usr/bin/ruby2.0 /usr/bin/ruby && sudo ln -sf /usr/bin/gem2.0 /usr/bin/gem
+
+  if ! gem list | grep bundler; then
+    sudo gem install bundler --no-ri --no-rdoc
+  fi
+
+  # Don't fail because we haven't added github.com's ssh key to our known_hosts:
+  cat <<EOF | sudo tee -a /etc/ssh/ssh_config > /dev/null
+Host github.com
+    StrictHostKeyChecking no
+EOF
+
+  sudo gem install eye --no-ri --no-rdoc
+}
+
+setup_fluentd() {
+  if [ -f /etc/init.d/td-agent ]; then
+    echo "fluentd is already installed"
+  else
+    curl -L http://toolbelt.treasuredata.com/sh/install-ubuntu-trusty-td-agent2.sh | sudo sh
+  fi
+
+  if td-agent-gem list | grep fluent-plugin-elasticsearch; then
+    echo "fluent-plugin-elasticsearch is already installed"
+  else
+    sudo td-agent-gem install fluent-plugin-elasticsearch -v 0.7.0 --no-ri --no-rdoc
+  fi
+
+  if grep "distribot" /etc/rsyslog.conf; then
+    echo "syslogd already forwarding events to fluentd"
+  else
+    echo '!distribot' | sudo tee -a /etc/rsyslog.conf
+    echo "*.* @127.0.0.1:42185" | sudo tee -a /etc/rsyslog.conf
+  fi
+
+  sudo cp provision/templates/fluentd.conf /etc/td-agent/td-agent.conf
+
+  sudo service rsyslog restart
+
+  # Finally, restart td-agent:
+  sudo service td-agent restart
+}
+
+setup_dependencies
+
+while ! ( echo -e "443\n6379\n5672" | xargs -i nc -w 1 -zv $INFRA_HOST {} ) ; do
+  echo "Waiting for infra to come up..."
+  sleep 5
+done
+
+setup_fluentd
+
+cd /var/www/distribot
+bundle
+
+echo '--------------------------------------------------------------------------'
+echo '--------------------------------------------------------------------------'
+echo '--------------------------------------------------------------------------'
+echo '--------------------------------------------------------------------------'
+echo '--------------------------------------------------------------------------'

data/provision/templates/fluentd.conf

@@ -0,0 +1,27 @@
+
+
+<source>
+  type syslog
+  port 42185
+  bind 127.0.0.1
+  tag syslog
+</source>
+
+<match **>
+  type copy
+  <store>
+    type stdout
+  </store>
+
+  <store>
+    type elasticsearch
+    logstash_format true
+    flush_interval 5s
+    type_name distribot
+
+    # We post data to https://<elasticsearch_hostname>/elastic/<index_name>/
+    hosts https://infra/elastic/
+    index_name distribot
+  </store>
+
+</match>

data/spec/distribot/bunny_connector_spec.rb

@@ -0,0 +1,196 @@
+require 'spec_helper'
+
+describe Distribot::BunnyConnector do
+
+  describe '#initialize(connection_args={})' do
+    before do
+      @amqp_url = 'amqp://distribot:distribot@172.17.0.2:5672'
+      expect(Bunny).to receive(:new).with(@amqp_url) do
+        bunny = double('bunny')
+        expect(bunny).to receive(:start).ordered
+        bunny
+      end
+    end
+
+    it 'initializes a new connector' do
+      connector = described_class.new(@amqp_url)
+      expect(connector).to be_a Distribot::BunnyConnector
+    end
+
+    it 'initializes subscribers as an empty array' do
+      connector = described_class.new(@amqp_url)
+      expect(connector.subscribers).to eq [ ]
+    end
+  end
+
+  describe '#channel' do
+    before do
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new
+    end
+    context 'the first time' do
+      it 'creates a new channel and returns it' do
+        bunny = double('bunny')
+        expect(bunny).to receive(:create_channel){ 'a-channel' }
+        expect(@connector).to receive(:bunny){ bunny }
+        expect(@connector.channel).to eq 'a-channel'
+      end
+    end
+    context 'after the first time' do
+      it 'returns the same first channel' do
+        bunny = double('bunny')
+        expect(bunny).to receive(:create_channel).exactly(1).times{ 'a-channel' }
+        expect(@connector).to receive(:bunny){ bunny }
+        expect(@connector.channel).to eq 'a-channel'
+        expect(@connector.channel).to eq 'a-channel'
+      end
+    end
+  end
+
+  describe '#queue_exists?(topic)' do
+    before do
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new(nil)
+      @topic = SecureRandom.uuid
+      @bunny = double('bunny')
+      expect(@connector).to receive(:bunny){ @bunny }
+    end
+    context 'when the queue' do
+      context 'exists' do
+        before do
+          expect(@bunny).to receive(:queue_exists?).with(@topic){ true }
+        end
+        it 'returns true' do
+          expect(@connector.queue_exists?(@topic)).to be_truthy
+        end
+      end
+      context 'does not exist' do
+        before do
+          expect(@bunny).to receive(:queue_exists?).with(@topic){ false }
+        end
+        it 'returns false' do
+          expect(@connector.queue_exists?(@topic)).to be_falsey
+        end
+      end
+    end
+  end
+
+  describe '#subscribe(topic, options={}, &block)' do
+    before do
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new(nil)
+      @topic = SecureRandom.uuid
+      expect_any_instance_of(Distribot::Subscription).to receive(:start).with(@topic, {}) do |&block|
+        block.call( id: 'hello' )
+      end
+    end
+    context 'when options[:solo]' do
+      context 'is truthy' do
+        it 'first calls setup, then subscribes to the topic, and calls the block when a message is received' do
+          expect(@connector).to receive(:setup)
+          @connector.subscribe(@topic, solo: true) do |msg|
+            @id = msg[:id]
+          end
+          expect(@id).to eq 'hello'
+        end
+      end
+      context 'is falsey' do
+        it 'subscribes to the topic, and calls the block when a message is received' do
+          @connector.subscribe(@topic) do |msg|
+            @id = msg[:id]
+          end
+          expect(@id).to eq 'hello'
+        end
+      end
+    end
+  end
+
+  describe '#subscribe_multi(topic, options={}, &block)' do
+    before do
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new(nil)
+      @topic = SecureRandom.uuid
+      expect_any_instance_of(Distribot::MultiSubscription).to receive(:start).with(@topic, {}) do |&block|
+        block.call( id: 'hello' )
+      end
+    end
+    it 'subscribes to the topic, and calls the block when a message is received' do
+      @connector.subscribe_multi(@topic) do |msg|
+        @id = msg[:id]
+      end
+      expect(@id).to eq 'hello'
+    end
+  end
+
+  describe '#publish(topic, message)' do
+    before do
+      @topic = SecureRandom.uuid
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new
+      @channel = double('channel')
+      expect(@channel).to receive(:queue).with(@topic, auto_delete: false, durable: true) do
+        queue = double('queue')
+        expect(queue).to receive(:name){ @topic }
+        queue
+      end
+      expect(@channel).to receive(:default_exchange) do
+        exchange = double('exchange')
+        expect(exchange).to receive(:publish).with( '{"hello":"world"}', routing_key: @topic)
+        exchange
+      end
+      expect(@connector).to receive(:channel).exactly(2).times{ @channel }
+    end
+    it 'publishes the message' do
+      @connector.publish(@topic, {hello: :world})
+    end
+  end
+
+  describe '#broadcast(topic, message)' do
+    before do
+      @topic = SecureRandom.uuid
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new
+      @channel = double('channel')
+      expect(@channel).to receive(:fanout).with(@topic) do
+        exchange = double('exchange')
+        expect(exchange).to receive(:publish).with('{"hello":"world"}', routing_key: @topic)
+        exchange
+      end
+      expect(@connector).to receive(:channel){ @channel }
+    end
+    it 'broadcasts a message on a fanout exchange' do
+      @connector.broadcast(@topic, {hello: :world})
+    end
+  end
+
+  describe '#stubbornly' do
+    context 'when the block' do
+      before do
+        expect_any_instance_of(described_class).to receive(:setup)
+      end
+      context 'raises an error' do
+        it 'keeps trying forever, until it stops raising an error' do
+          @return_value = SecureRandom.uuid
+          thing = described_class.new
+          @max_tries = 3
+          @total_tries = 0
+          expect(thing.send(:stubbornly, :foo){
+            if @total_tries >= @max_tries
+              @return_value
+            else
+              @total_tries += 1
+              raise Timeout::Error.new
+            end
+          }).to eq @return_value
+        end
+      end
+      context 'does not raise an error' do
+        it 'returns the result of the block' do
+          @return_value = SecureRandom.uuid
+          thing = described_class.new
+          expect(thing.send(:stubbornly, :foo){ @return_value }).to eq @return_value
+        end
+      end
+    end
+  end
+end

data/spec/distribot/connection_sharer_spec.rb

@@ -0,0 +1,34 @@
+require 'spec_helper'
+
+describe Distribot::ConnectionSharer do
+  describe '#initialize(bunny)' do
+    before do
+      @bunny = SecureRandom.uuid
+      @sharer = described_class.new(@bunny)
+    end
+    it 'sets self.bunny' do
+      expect(@sharer.bunny).to eq @bunny
+    end
+  end
+
+  describe '#channel' do
+    before do
+      @bunny = double('bunny')
+      expect(@bunny).to receive(:create_channel){ SecureRandom.uuid }
+      @sharer = described_class.new(@bunny)
+    end
+    context 'the first time' do
+      it 'creates a new channel and stores it' do
+        expect(SecureRandom).to receive(:uuid){ 'your-channel' }
+        expect(@sharer.channel).to eq 'your-channel'
+      end
+    end
+    context 'each subsequent time' do
+      it 'returns the original channel' do
+        expect(SecureRandom).to receive(:uuid).exactly(1).times.and_call_original
+        first_channel = @sharer.channel
+        expect(@sharer.channel).to eq first_channel
+      end
+    end
+  end
+end

data/spec/distribot/connector_spec.rb

@@ -0,0 +1,63 @@
+
+require 'spec_helper'
+
+describe Distribot::Connector do
+
+  describe '#initialize(connection_args={})' do
+    before do
+      @amqp_url = 'amqp://distribot:distribot@172.17.0.2:5672'
+      expect(Bunny).to receive(:new).with(@amqp_url) do
+        bunny = double('bunny')
+        expect(bunny).to receive(:start).ordered
+        expect(bunny).to receive(:create_channel).ordered do
+          channel = double('channel')
+          expect(channel).to receive(:prefetch).with(1)
+          channel
+        end
+        bunny
+      end
+    end
+
+    it 'initializes a new connector' do
+      connector = described_class.new(@amqp_url)
+      expect(connector).to be_a Distribot::Connector
+    end
+  end
+
+  describe '#queues' do
+    before do
+      expect_any_instance_of(described_class).to receive(:setup)
+      @connector = described_class.new
+      @queues = %w(
+        distribot.flow.created
+        distribot.flow.finished
+        distribot.flow.handler.CheapWorker.1.0.0.enumerate
+        distribot.flow.handler.CheapWorker.1.0.0.tasks
+        distribot.flow.handler.FastWorker.1.0.0.enumerate
+        distribot.flow.handler.FastWorker.1.0.0.tasks
+        distribot.flow.handler.ForeignWorker.1.0.0.enumerate
+        distribot.flow.handler.ForeignWorker.1.0.0.tasks
+        distribot.flow.handler.GoodWorker.1.0.0.enumerate
+        distribot.flow.handler.GoodWorker.1.0.0.tasks
+        distribot.flow.handler.HardWorker.1.0.0.enumerate
+        distribot.flow.handler.HardWorker.1.0.0.tasks
+        distribot.flow.handler.SlowWorker.1.0.0.enumerate
+        distribot.flow.handler.SlowWorker.1.0.0.tasks
+        distribot.flow.handler.finished
+        distribot.flow.phase.finished
+        distribot.flow.phase.started
+        distribot.flow.task.finished
+      )
+      @queues_json = @queues.to_a.map{|name| {name: name} }.to_json
+      Wrest.logger = Logger.new('/dev/null')
+      stub_request(:get, "http://localhost:15672/api/queues").
+        with(:headers => {'Accept'=>'*/*', 'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3', 'User-Agent'=>'Ruby'}).
+        to_return(:status => 200, :body => @queues_json, :headers => {'Content-Type' => 'application/json'})
+    end
+    it 'returns all the queues from /api/queues on rabbitmq' do
+      result = @connector.queues
+      expect(result).to eq(@queues)
+    end
+  end
+
+end

data/spec/distribot/flow_created_handler_spec.rb

@@ -0,0 +1,32 @@
+require 'spec_helper'
+
+describe Distribot::FlowCreatedHandler do
+  describe 'definition' do
+    it 'subscribes to the correct queue' do
+      expect(Distribot::Handler.queue_for(described_class)).to eq 'distribot.flow.created'
+    end
+    it 'declares a valid handler' do
+      expect(Distribot::Handler.handler_for(described_class)).to eq :callback
+    end
+    it 'has a method matching the handler name' do
+      expect(Distribot).to receive(:subscribe)
+      expect(described_class.new).to respond_to :callback
+    end
+  end
+
+  describe '#callback' do
+    before do
+      @flow_id = SecureRandom.uuid
+      expect(Distribot::Flow).to receive(:find).with(@flow_id) do
+        flow = double('flow')
+        expect(flow).to receive(:next_phase){'phase2'}
+        expect(flow).to receive(:transition_to!).with('phase2'){ true }
+        flow
+      end
+      expect(Distribot).to receive(:subscribe)
+    end
+    it 'transitions to the next phase' do
+      expect(described_class.new.callback(flow_id: @flow_id)).to be_truthy
+    end
+  end
+end

data/spec/distribot/flow_finished_handler_spec.rb

@@ -0,0 +1,32 @@
+require 'spec_helper'
+
+describe Distribot::FlowFinishedHandler do
+  describe 'definition' do
+    it 'subscribes to the distribot.flow.finished queue' do
+      expect(Distribot::Handler.queue_for(described_class)).to eq 'distribot.flow.finished'
+    end
+    it 'declares a valid handler' do
+      expect(Distribot::Handler.handler_for(described_class)).to eq :callback
+    end
+    it 'has a method matching the handler name' do
+      expect(Distribot).to receive(:subscribe)
+      expect(described_class.new).to respond_to :callback
+    end
+  end
+
+  describe '#callback' do
+    before do
+      @message = {
+        flow_id: 'xxx'
+      }
+      redis = double('redis')
+      expect(redis).to receive(:decr).with('distribot.flows.running')
+      expect(redis).to receive(:srem).with('distribot.flows.active', @message[:flow_id])
+      expect(Distribot).to receive(:redis).exactly(2).times{ redis }
+      expect(Distribot).to receive(:subscribe)
+    end
+    it 'decrements the running tally of how many flows are currently in process' do
+      described_class.new.callback(@message)
+    end
+  end
+end