net-ssh-multi 1.0.0

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,546 @@
1
+ require 'thread'
2
+ require 'net/ssh/gateway'
3
+ require 'net/ssh/multi/server'
4
+ require 'net/ssh/multi/dynamic_server'
5
+ require 'net/ssh/multi/server_list'
6
+ require 'net/ssh/multi/channel'
7
+ require 'net/ssh/multi/pending_connection'
8
+ require 'net/ssh/multi/session_actions'
9
+ require 'net/ssh/multi/subsession'
10
+
11
+ module Net; module SSH; module Multi
12
+ # Represents a collection of connections to various servers. It provides an
13
+ # interface for organizing the connections (#group), as well as a way to
14
+ # scope commands to a subset of all connections (#with). You can also provide
15
+ # a default gateway connection that servers should use when connecting
16
+ # (#via). It exposes an interface similar to Net::SSH::Connection::Session
17
+ # for opening SSH channels and executing commands, allowing for these
18
+ # operations to be done in parallel across multiple connections.
19
+ #
20
+ # Net::SSH::Multi.start do |session|
21
+ # # access servers via a gateway
22
+ # session.via 'gateway', 'gateway-user'
23
+ #
24
+ # # define the servers we want to use
25
+ # session.use 'user1@host1'
26
+ # session.use 'user2@host2'
27
+ #
28
+ # # define servers in groups for more granular access
29
+ # session.group :app do
30
+ # session.use 'user@app1'
31
+ # session.use 'user@app2'
32
+ # end
33
+ #
34
+ # # execute commands on all servers
35
+ # session.exec "uptime"
36
+ #
37
+ # # execute commands on a subset of servers
38
+ # session.with(:app).exec "hostname"
39
+ #
40
+ # # run the aggregated event loop
41
+ # session.loop
42
+ # end
43
+ #
44
+ # Note that connections are established lazily, as soon as they are needed.
45
+ # You can force the connections to be opened immediately, though, using the
46
+ # #connect! method.
47
+ #
48
+ # == Concurrent Connection Limiting
49
+ #
50
+ # Sometimes you may be dealing with a large number of servers, and if you
51
+ # try to have connections open to all of them simultaneously you'll run into
52
+ # open file handle limitations and such. If this happens to you, you can set
53
+ # the #concurrent_connections property of the session. Net::SSH::Multi will
54
+ # then ensure that no more than this number of connections are ever open
55
+ # simultaneously.
56
+ #
57
+ # Net::SSH::Multi.start(:concurrent_connections => 5) do |session|
58
+ # # ...
59
+ # end
60
+ #
61
+ # Opening channels and executing commands will still work exactly as before,
62
+ # but Net::SSH::Multi will transparently close finished connections and open
63
+ # pending ones.
64
+ #
65
+ # == Controlling Connection Errors
66
+ #
67
+ # By default, Net::SSH::Multi will raise an exception if a connection error
68
+ # occurs when connecting to a server. This will typically bubble up and abort
69
+ # the entire connection process. Sometimes, however, you might wish to ignore
70
+ # connection errors, for instance when starting a daemon on a large number of
71
+ # boxes and you know that some of the boxes are going to be unavailable.
72
+ #
73
+ # To do this, simply set the #on_error property of the session to :ignore
74
+ # (or to :warn, if you want a warning message when a connection attempt
75
+ # fails):
76
+ #
77
+ # Net::SSH::Multi.start(:on_error => :ignore) do |session|
78
+ # # ...
79
+ # end
80
+ #
81
+ # The default is :fail, which causes the exception to bubble up. Additionally,
82
+ # you can specify a Proc object as the value for #on_error, which will be
83
+ # invoked with the server in question if the connection attempt fails. You
84
+ # can force the connection attempt to retry by throwing the :go symbol, with
85
+ # :retry as the payload, or force the exception to be reraised by throwing
86
+ # :go with :raise as the payload:
87
+ #
88
+ # handler = Proc.new do |server|
89
+ # server[:connection_attempts] ||= 0
90
+ # if server[:connection_attempts] < 3
91
+ # server[:connection_attempts] += 1
92
+ # throw :go, :retry
93
+ # else
94
+ # throw :go, :raise
95
+ # end
96
+ # end
97
+ #
98
+ # Net::SSH::Multi.start(:on_error => handler) do |session|
99
+ # # ...
100
+ # end
101
+ #
102
+ # Any other thrown value (or no thrown value at all) will result in the
103
+ # failure being ignored.
104
+ #
105
+ # == Lazily Evaluated Server Definitions
106
+ #
107
+ # Sometimes you might be dealing with an environment where you don't know the
108
+ # names or addresses of the servers until runtime. You can certainly dynamically
109
+ # build server names and pass them to #use, but if the operation to determine
110
+ # the server names is expensive, you might want to defer it until the server
111
+ # is actually needed (especially if the logic of your program is such that
112
+ # you might not even need to connect to that server every time the program
113
+ # runs).
114
+ #
115
+ # You can do this by passing a block to #use:
116
+ #
117
+ # session.use do |opt|
118
+ # lookup_ip_address_of_remote_host
119
+ # end
120
+ #
121
+ # See #use for more information about this usage.
122
+ class Session
123
+ include SessionActions
124
+
125
+ # The Net::SSH::Multi::ServerList managed by this session.
126
+ attr_reader :server_list
127
+
128
+ # The default Net::SSH::Gateway instance to use to connect to the servers.
129
+ # If +nil+, no default gateway will be used.
130
+ attr_reader :default_gateway
131
+
132
+ # The hash of group definitions, mapping each group name to a corresponding
133
+ # Net::SSH::Multi::ServerList.
134
+ attr_reader :groups
135
+
136
+ # The number of allowed concurrent connections. No more than this number
137
+ # of sessions will be open at any given time.
138
+ attr_accessor :concurrent_connections
139
+
140
+ # How connection errors should be handled. This defaults to :fail, but
141
+ # may be set to :ignore if connection errors should be ignored, or
142
+ # :warn if connection errors should cause a warning.
143
+ attr_accessor :on_error
144
+
145
+ # The default user name to use when connecting to a server. If a user name
146
+ # is not given for a particular server, this value will be used. It defaults
147
+ # to ENV['USER'] || ENV['USERNAME'], or "unknown" if neither of those are
148
+ # set.
149
+ attr_accessor :default_user
150
+
151
+ # The number of connections that are currently open.
152
+ attr_reader :open_connections #:nodoc:
153
+
154
+ # The list of "open" groups, which will receive subsequent server definitions.
155
+ # See #use and #group.
156
+ attr_reader :open_groups #:nodoc:
157
+
158
+ # Creates a new Net::SSH::Multi::Session instance. Initially, it contains
159
+ # no server definitions, no group definitions, and no default gateway.
160
+ #
161
+ # You can set the #concurrent_connections property in the options. Setting
162
+ # it to +nil+ (the default) will cause Net::SSH::Multi to ignore any
163
+ # concurrent connection limit and allow all defined sessions to be open
164
+ # simultaneously. Setting it to an integer will cause Net::SSH::Multi to
165
+ # allow no more than that number of concurrently open sessions, opening
166
+ # subsequent sessions only when other sessions finish and close.
167
+ #
168
+ # Net::SSH::Multi.start(:concurrent_connections => 10) do |session|
169
+ # session.use ...
170
+ # end
171
+ def initialize(options={})
172
+ @server_list = ServerList.new
173
+ @groups = Hash.new { |h,k| h[k] = ServerList.new }
174
+ @gateway = nil
175
+ @open_groups = []
176
+ @connect_threads = []
177
+ @on_error = :fail
178
+ @default_user = ENV['USER'] || ENV['USERNAME'] || "unknown"
179
+
180
+ @open_connections = 0
181
+ @pending_sessions = []
182
+ @session_mutex = Mutex.new
183
+
184
+ options.each { |opt, value| send("#{opt}=", value) }
185
+ end
186
+
187
+ # At its simplest, this associates a named group with a server definition.
188
+ # It can be used in either of two ways:
189
+ #
190
+ # First, you can use it to associate a group (or array of groups) with a
191
+ # server definition (or array of server definitions). The server definitions
192
+ # must already exist in the #server_list array (typically by calling #use):
193
+ #
194
+ # server1 = session.use('host1', 'user1')
195
+ # server2 = session.use('host2', 'user2')
196
+ # session.group :app => server1, :web => server2
197
+ # session.group :staging => [server1, server2]
198
+ # session.group %w(xen linux) => server2
199
+ # session.group %w(rackspace backup) => [server1, server2]
200
+ #
201
+ # Secondly, instead of a mapping of groups to servers, you can just
202
+ # provide a list of group names, and then a block. Inside the block, any
203
+ # calls to #use will automatically associate the new server definition with
204
+ # those groups. You can nest #group calls, too, which will aggregate the
205
+ # group definitions.
206
+ #
207
+ # session.group :rackspace, :backup do
208
+ # session.use 'host1', 'user1'
209
+ # session.group :xen do
210
+ # session.use 'host2', 'user2'
211
+ # end
212
+ # end
213
+ def group(*args)
214
+ mapping = args.last.is_a?(Hash) ? args.pop : {}
215
+
216
+ if mapping.any? && block_given?
217
+ raise ArgumentError, "must provide group mapping OR block, not both"
218
+ elsif block_given?
219
+ begin
220
+ saved_groups = open_groups.dup
221
+ open_groups.concat(args.map { |a| a.to_sym }).uniq!
222
+ yield self if block_given?
223
+ ensure
224
+ open_groups.replace(saved_groups)
225
+ end
226
+ else
227
+ mapping.each do |key, value|
228
+ (open_groups + Array(key)).uniq.each do |grp|
229
+ groups[grp.to_sym].concat(Array(value))
230
+ end
231
+ end
232
+ end
233
+ end
234
+
235
+ # Sets up a default gateway to use when establishing connections to servers.
236
+ # Note that any servers defined prior to this invocation will not use the
237
+ # default gateway; it only affects servers defined subsequently.
238
+ #
239
+ # session.via 'gateway.host', 'user'
240
+ #
241
+ # You may override the default gateway on a per-server basis by passing the
242
+ # :via key to the #use method; see #use for details.
243
+ def via(host, user, options={})
244
+ @default_gateway = Net::SSH::Gateway.new(host, user, options)
245
+ self
246
+ end
247
+
248
+ # Defines a new server definition, to be managed by this session. The
249
+ # server is at the given +host+, and will be connected to as the given
250
+ # +user+. The other options are passed as-is to the Net::SSH session
251
+ # constructor.
252
+ #
253
+ # If a default gateway has been specified previously (with #via) it will
254
+ # be passed to the new server definition. You can override this by passing
255
+ # a different Net::SSH::Gateway instance (or +nil+) with the :via key in
256
+ # the +options+.
257
+ #
258
+ # session.use 'host'
259
+ # session.use 'user@host2', :via => nil
260
+ # session.use 'host3', :user => "user3", :via => Net::SSH::Gateway.new('gateway.host', 'user')
261
+ #
262
+ # If only a single host is given, the new server instance is returned. You
263
+ # can give multiple hosts at a time, though, in which case an array of
264
+ # server instances will be returned.
265
+ #
266
+ # server1, server2 = session.use "host1", "host2"
267
+ #
268
+ # If given a block, this will save the block as a Net::SSH::Multi::DynamicServer
269
+ # definition, to be evaluated lazily the first time the server is needed.
270
+ # The block will recive any options hash given to #use, and should return
271
+ # +nil+ (if no servers are to be added), a String or an array of Strings
272
+ # (to be interpreted as a connection specification), or a Server or an
273
+ # array of Servers.
274
+ def use(*hosts, &block)
275
+ options = hosts.last.is_a?(Hash) ? hosts.pop : {}
276
+ options = { :via => default_gateway }.merge(options)
277
+
278
+ results = hosts.map do |host|
279
+ server_list.add(Server.new(self, host, options))
280
+ end
281
+
282
+ if block
283
+ results << server_list.add(DynamicServer.new(self, options, block))
284
+ end
285
+
286
+ group [] => results
287
+ results.length > 1 ? results : results.first
288
+ end
289
+
290
+ # Essentially an alias for #servers_for without any arguments. This is used
291
+ # primarily to satistfy the expectations of the Net::SSH::Multi::SessionActions
292
+ # module.
293
+ def servers
294
+ servers_for
295
+ end
296
+
297
+ # Returns the set of servers that match the given criteria. It can be used
298
+ # in any (or all) of three ways.
299
+ #
300
+ # First, you can omit any arguments. In this case, the full list of servers
301
+ # will be returned.
302
+ #
303
+ # all = session.servers_for
304
+ #
305
+ # Second, you can simply specify a list of group names. All servers in all
306
+ # named groups will be returned. If a server belongs to multiple matching
307
+ # groups, then it will appear only once in the list (the resulting list
308
+ # will contain only unique servers).
309
+ #
310
+ # servers = session.servers_for(:app, :db)
311
+ #
312
+ # Last, you can specify a hash with group names as keys, and property
313
+ # constraints as the values. These property constraints are either "only"
314
+ # constraints (which restrict the set of servers to "only" those that match
315
+ # the given properties) or "except" constraints (which restrict the set of
316
+ # servers to those whose properties do _not_ match). Properties are described
317
+ # when the server is defined (via the :properties key):
318
+ #
319
+ # session.group :db do
320
+ # session.use 'dbmain', 'user', :properties => { :primary => true }
321
+ # session.use 'dbslave', 'user2'
322
+ # session.use 'dbslve2', 'user2'
323
+ # end
324
+ #
325
+ # # return ONLY on the servers in the :db group which have the :primary
326
+ # # property set to true.
327
+ # primary = session.servers_for(:db => { :only => { :primary => true } })
328
+ #
329
+ # You can, naturally, combine these methods:
330
+ #
331
+ # # all servers in :app and :web, and all servers in :db with the
332
+ # # :primary property set to true
333
+ # servers = session.servers_for(:app, :web, :db => { :only => { :primary => true } })
334
+ def servers_for(*criteria)
335
+ if criteria.empty?
336
+ server_list.flatten
337
+ else
338
+ # normalize the criteria list, so that every entry is a key to a
339
+ # criteria hash (possibly empty).
340
+ criteria = criteria.inject({}) do |hash, entry|
341
+ case entry
342
+ when Hash then hash.merge(entry)
343
+ else hash.merge(entry => {})
344
+ end
345
+ end
346
+
347
+ list = criteria.inject([]) do |aggregator, (group, properties)|
348
+ raise ArgumentError, "the value for any group must be a Hash, but got a #{properties.class} for #{group.inspect}" unless properties.is_a?(Hash)
349
+ bad_keys = properties.keys - [:only, :except]
350
+ raise ArgumentError, "unknown constraint(s) #{bad_keys.inspect} for #{group.inspect}" unless bad_keys.empty?
351
+
352
+ servers = groups[group].select do |server|
353
+ (properties[:only] || {}).all? { |prop, value| server[prop] == value } &&
354
+ !(properties[:except] || {}).any? { |prop, value| server[prop] == value }
355
+ end
356
+
357
+ aggregator.concat(servers)
358
+ end
359
+
360
+ list.uniq
361
+ end
362
+ end
363
+
364
+ # Returns a new Net::SSH::Multi::Subsession instance consisting of the
365
+ # servers that meet the given criteria. If a block is given, the
366
+ # subsession will be yielded to it. See #servers_for for a discussion of
367
+ # how these criteria are interpreted.
368
+ #
369
+ # session.with(:app).exec('hostname')
370
+ #
371
+ # session.with(:app, :db => { :primary => true }) do |s|
372
+ # s.exec 'date'
373
+ # s.exec 'uptime'
374
+ # end
375
+ def with(*groups)
376
+ subsession = Subsession.new(self, servers_for(*groups))
377
+ yield subsession if block_given?
378
+ subsession
379
+ end
380
+
381
+ # Works as #with, but for specific servers rather than groups. It will
382
+ # return a new subsession (Net::SSH::Multi::Subsession) consisting of
383
+ # the given servers. (Note that it requires that the servers in question
384
+ # have been created via calls to #use on this session object, or things
385
+ # will not work quite right.) If a block is given, the new subsession
386
+ # will also be yielded to the block.
387
+ #
388
+ # srv1 = session.use('host1', 'user')
389
+ # srv2 = session.use('host2', 'user')
390
+ # # ...
391
+ # session.on(srv1, srv2).exec('hostname')
392
+ def on(*servers)
393
+ subsession = Subsession.new(self, servers)
394
+ yield subsession if block_given?
395
+ subsession
396
+ end
397
+
398
+ # Closes the multi-session by shutting down all open server sessions, and
399
+ # the default gateway (if one was specified using #via). Note that other
400
+ # gateway connections (e.g., those passed to #use directly) will _not_ be
401
+ # closed by this method, and must be managed externally.
402
+ def close
403
+ server_list.each { |server| server.close_channels }
404
+ loop(0) { busy?(true) }
405
+ server_list.each { |server| server.close }
406
+ default_gateway.shutdown! if default_gateway
407
+ end
408
+
409
+ alias :loop_forever :loop
410
+
411
+ # Run the aggregated event loop for all open server sessions, until the given
412
+ # block returns +false+. If no block is given, the loop will run for as
413
+ # long as #busy? returns +true+ (in other words, for as long as there are
414
+ # any (non-invisible) channels open).
415
+ def loop(wait=nil, &block)
416
+ running = block || Proc.new { |c| busy? }
417
+ loop_forever { break unless process(wait, &running) }
418
+ end
419
+
420
+ # Run a single iteration of the aggregated event loop for all open server
421
+ # sessions. The +wait+ parameter indicates how long to wait for an event
422
+ # to appear on any of the different sessions; +nil+ (the default) means
423
+ # "wait forever". If the block is given, then it will be used to determine
424
+ # whether #process returns +true+ (the block did not return +false+), or
425
+ # +false+ (the block returned +false+).
426
+ def process(wait=nil, &block)
427
+ realize_pending_connections!
428
+ wait = @connect_threads.any? ? 0 : wait
429
+
430
+ return false unless preprocess(&block)
431
+
432
+ readers = server_list.map { |s| s.readers }.flatten
433
+ writers = server_list.map { |s| s.writers }.flatten
434
+
435
+ readers, writers, = IO.select(readers, writers, nil, wait)
436
+
437
+ if readers
438
+ return postprocess(readers, writers)
439
+ else
440
+ return true
441
+ end
442
+ end
443
+
444
+ # Runs the preprocess stage on all servers. Returns false if the block
445
+ # returns false, and true if there either is no block, or it returns true.
446
+ # This is called as part of the #process method.
447
+ def preprocess(&block) #:nodoc:
448
+ return false if block && !block[self]
449
+ server_list.each { |server| server.preprocess }
450
+ block.nil? || block[self]
451
+ end
452
+
453
+ # Runs the postprocess stage on all servers. Always returns true. This is
454
+ # called as part of the #process method.
455
+ def postprocess(readers, writers) #:nodoc:
456
+ server_list.each { |server| server.postprocess(readers, writers) }
457
+ true
458
+ end
459
+
460
+ # Takes the #concurrent_connections property into account, and tries to
461
+ # return a new session for the given server. If the concurrent connections
462
+ # limit has been reached, then a Net::SSH::Multi::PendingConnection instance
463
+ # will be returned instead, which will be realized into an actual session
464
+ # as soon as a slot opens up.
465
+ #
466
+ # If +force+ is true, the concurrent_connections check is skipped and a real
467
+ # connection is always returned.
468
+ def next_session(server, force=false) #:nodoc:
469
+ # don't retry a failed attempt
470
+ return nil if server.failed?
471
+
472
+ @session_mutex.synchronize do
473
+ if !force && concurrent_connections && concurrent_connections <= open_connections
474
+ connection = PendingConnection.new(server)
475
+ @pending_sessions << connection
476
+ return connection
477
+ end
478
+
479
+ @open_connections += 1
480
+ end
481
+
482
+ begin
483
+ server.new_session
484
+ rescue Exception => e
485
+ server.fail!
486
+ @session_mutex.synchronize { @open_connections -= 1 }
487
+
488
+ case on_error
489
+ when :ignore then
490
+ # do nothing
491
+ when :warn then
492
+ warn("error connecting to #{server}: #{e.class} (#{e.message})")
493
+ when Proc then
494
+ go = catch(:go) { on_error.call(server); nil }
495
+ case go
496
+ when nil, :ignore then # nothing
497
+ when :retry then retry
498
+ when :raise then raise
499
+ else warn "unknown 'go' command: #{go.inspect}"
500
+ end
501
+ else
502
+ raise
503
+ end
504
+
505
+ return nil
506
+ end
507
+ end
508
+
509
+ # Tells the session that the given server has closed its connection. The
510
+ # session indicates that a new connection slot is available, which may be
511
+ # filled by the next pending connection on the next event loop iteration.
512
+ def server_closed(server) #:nodoc:
513
+ @session_mutex.synchronize do
514
+ unless @pending_sessions.delete(server.session)
515
+ @open_connections -= 1
516
+ end
517
+ end
518
+ end
519
+
520
+ # Invoked by the event loop. If there is a concurrent_connections limit in
521
+ # effect, this will close any non-busy sessions and try to open as many
522
+ # new sessions as it can. It does this in threads, so that existing processing
523
+ # can continue.
524
+ #
525
+ # If there is no concurrent_connections limit in effect, then this method
526
+ # does nothing.
527
+ def realize_pending_connections! #:nodoc:
528
+ return unless concurrent_connections
529
+
530
+ server_list.each do |server|
531
+ server.close if !server.busy?(true)
532
+ server.update_session!
533
+ end
534
+
535
+ @connect_threads.delete_if { |t| !t.alive? }
536
+
537
+ count = concurrent_connections ? (concurrent_connections - open_connections) : @pending_sessions.length
538
+ count.times do
539
+ session = @pending_sessions.pop or break
540
+ @connect_threads << Thread.new do
541
+ session.replace_with(next_session(session.server, true))
542
+ end
543
+ end
544
+ end
545
+ end
546
+ end; end; end