Class: Procodile::TCPProxy

Inherits:
Object
  • Object
show all
Defined in:
lib/procodile/tcp_proxy.rb

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(supervisor) ⇒ TCPProxy

Returns a new instance of TCPProxy.



10
11
12
13
14
15
16
 
# File 'lib/procodile/tcp_proxy.rb', line 10

def initialize(supervisor)
  @supervisor = supervisor
  @thread = nil
  @listeners = {}
  @stopped_processes = []
  @sp_reader, @sp_writer = IO.pipe
end

Class Method Details

.start(supervisor) ⇒ Object 



4
5
6
7
8
 
# File 'lib/procodile/tcp_proxy.rb', line 4

def self.start(supervisor)
  proxy = new(supervisor)
  proxy.start
  proxy
end

Instance Method Details

#add_process(process) ⇒ Object 



26
27
28
29
30
31
32
33
34
35
 
# File 'lib/procodile/tcp_proxy.rb', line 26

def add_process(process)
  if process.proxy?
    @listeners[TCPServer.new(process.proxy_address, process.proxy_port)] = process
    Procodile.log nil, 'proxy', "Proxying traffic on #{process.proxy_address}:#{process.proxy_port} to #{process.name}".color(32)
    @sp_writer.write_nonblock('.')
  end
rescue => e
  Procodile.log nil, 'proxy', "Exception: #{e.class}: #{e.message}"
  Procodile.log nil, 'proxy', e.backtrace[0,5].join("\n")
end

#handle_client(client, server) ⇒ Object 



72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
 
# File 'lib/procodile/tcp_proxy.rb', line 72

def handle_client(client, server)
  process = @listeners[server]
  instances = @supervisor.processes[process] || []
  if instances.empty?
    Procodile.log nil, 'proxy', "There are no processes running for #{process.name}"
  else
    instance = instances[rand(instances.size)]
    backend_socket = TCPSocket.new('127.0.0.1', instance.port) rescue nil
    if backend_socket.nil?
      Procodile.log nil, 'proxy', "Could not connect to #{instance.description}:#{instance.port}"
      return
    end
    readers = {:backend => backend_socket, :client => client}
    loop do
      io = IO.select(readers.values, nil, nil, 0.5)
      if io && io.first
        io.first.each do |io|
          readers.keys.each do |key|
            next unless readers[key] == io
            opposite_side = key == :client ? :backend : :client
            if io.eof?
              readers[opposite_side].shutdown(Socket::SHUT_WR) rescue nil
              readers.delete(opposite_side)
            else
              readers[opposite_side].write(io.readpartial(1024)) rescue nil
            end
          end
        end
      end
    end
  end
rescue => e
  Procodile.log nil, 'proxy', "Exception: #{e.class}: #{e.message}"
  Procodile.log nil, 'proxy', e.backtrace[0,5].join("\n")
ensure
  backend_socket.close rescue nil
  client.close rescue nil
end

#listen ⇒ Object 



42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
 
# File 'lib/procodile/tcp_proxy.rb', line 42

def listen
  loop do
    io = IO.select([@sp_reader] + @listeners.keys, nil, nil, 30)
    if io && io.first
      io.first.each do |io|
        if io == @sp_reader
          io.read_nonblock(999)
          next
        end

        Thread.new(io.accept, io) do |client, server|
          handle_client(client, server)
        end
      end
    end

    @stopped_processes.reject do |process|
      if io = @listeners.key(process)
        Procodile.log nil, 'proxy', "Stopped proxy listener for #{process.name}"
        io.close
        @listeners.delete(io)
      end
      true
    end
  end
rescue => e
  Procodile.log nil, 'proxy', "Exception: #{e.class}: #{e.message}"
  Procodile.log nil, 'proxy', e.backtrace[0,5].join("\n")
end

#remove_process(process) ⇒ Object 



37
38
39
40
 
# File 'lib/procodile/tcp_proxy.rb', line 37

def remove_process(process)
  @stopped_processes << process
  @sp_writer.write_nonblock('.')
end

#start ⇒ Object 



18
19
20
21
22
23
24
 
# File 'lib/procodile/tcp_proxy.rb', line 18

def start
  @supervisor.config.processes.each { |_, p| add_process(p) }
  Thread.new do
    listen
    Procodile.log nil, 'proxy', "Stopped listening on all ports"
  end
end