switchman 2.0.11 → 3.0.4

Sign up to get free protection for your applications and to get access to all the features.
Files changed (56) hide show
  1. checksums.yaml +4 -4
  2. data/Rakefile +10 -2
  3. data/app/models/switchman/shard.rb +234 -270
  4. data/app/models/switchman/unsharded_record.rb +7 -0
  5. data/db/migrate/20130328212039_create_switchman_shards.rb +1 -1
  6. data/db/migrate/20130328224244_create_default_shard.rb +5 -5
  7. data/db/migrate/20161206323434_add_back_default_string_limits_switchman.rb +1 -0
  8. data/db/migrate/20180828183945_add_default_shard_index.rb +1 -1
  9. data/db/migrate/20180828192111_add_timestamps_to_shards.rb +7 -5
  10. data/db/migrate/20190114212900_add_unique_name_indexes.rb +5 -3
  11. data/lib/switchman.rb +3 -3
  12. data/lib/switchman/action_controller/caching.rb +2 -2
  13. data/lib/switchman/active_record/abstract_adapter.rb +1 -0
  14. data/lib/switchman/active_record/association.rb +78 -89
  15. data/lib/switchman/active_record/attribute_methods.rb +106 -52
  16. data/lib/switchman/active_record/base.rb +58 -59
  17. data/lib/switchman/active_record/calculations.rb +73 -66
  18. data/lib/switchman/active_record/connection_pool.rb +14 -41
  19. data/lib/switchman/active_record/database_configurations.rb +34 -0
  20. data/lib/switchman/active_record/database_configurations/database_config.rb +13 -0
  21. data/lib/switchman/active_record/finder_methods.rb +11 -16
  22. data/lib/switchman/active_record/log_subscriber.rb +4 -8
  23. data/lib/switchman/active_record/migration.rb +14 -43
  24. data/lib/switchman/active_record/model_schema.rb +1 -1
  25. data/lib/switchman/active_record/persistence.rb +4 -6
  26. data/lib/switchman/active_record/postgresql_adapter.rb +32 -160
  27. data/lib/switchman/active_record/predicate_builder.rb +1 -1
  28. data/lib/switchman/active_record/query_cache.rb +18 -19
  29. data/lib/switchman/active_record/query_methods.rb +179 -182
  30. data/lib/switchman/active_record/reflection.rb +6 -10
  31. data/lib/switchman/active_record/relation.rb +34 -29
  32. data/lib/switchman/active_record/spawn_methods.rb +27 -29
  33. data/lib/switchman/active_record/statement_cache.rb +18 -35
  34. data/lib/switchman/active_record/tasks/database_tasks.rb +16 -0
  35. data/lib/switchman/active_support/cache.rb +3 -5
  36. data/lib/switchman/arel.rb +13 -8
  37. data/lib/switchman/database_server.rb +122 -144
  38. data/lib/switchman/default_shard.rb +52 -16
  39. data/lib/switchman/engine.rb +61 -57
  40. data/lib/switchman/environment.rb +4 -8
  41. data/lib/switchman/errors.rb +1 -0
  42. data/lib/switchman/guard_rail.rb +6 -19
  43. data/lib/switchman/guard_rail/relation.rb +5 -7
  44. data/lib/switchman/r_spec_helper.rb +29 -37
  45. data/lib/switchman/rails.rb +14 -12
  46. data/lib/switchman/sharded_instrumenter.rb +1 -1
  47. data/lib/switchman/standard_error.rb +15 -3
  48. data/lib/switchman/test_helper.rb +6 -10
  49. data/lib/switchman/version.rb +1 -1
  50. data/lib/tasks/switchman.rake +54 -69
  51. metadata +85 -44
  52. data/lib/switchman/active_record/batches.rb +0 -11
  53. data/lib/switchman/active_record/connection_handler.rb +0 -172
  54. data/lib/switchman/active_record/where_clause_factory.rb +0 -36
  55. data/lib/switchman/connection_pool_proxy.rb +0 -173
  56. data/lib/switchman/schema_cache.rb +0 -28
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: aa20d460683c12b9e6528f0956c83dec5fadcc4247d551fbd84361757d5452a0
4
- data.tar.gz: 79a6a917b940d6e0bc5b10288f757d002635526e083f3c123698976fb5e96483
3
+ metadata.gz: 1ffc8c73f2ac38cf0440d1e556e1c7c2af4ad8d665e40893caf76ae228bc702f
4
+ data.tar.gz: a4567920bb1cff999d77845d0de9f204826cce7ef9da1c2164802230923b56af
5
5
  SHA512:
6
- metadata.gz: 55aca398595c2fda0e32f4aa35e10f42ded64d3154d1dd9bb5cee8f0a0682d61317f8486b156e2722ed3c381dd6cc42710729e19d859b3ade543991469464835
7
- data.tar.gz: 4b5bfb169f36aa3a69f716195316941f648d8b06ae3c9a9f09c5b5835dd67005ddbc7fe7cbe864b278aefeb3abdeaa8f65292f00dcd29c4cf1011c4e852bec64
6
+ metadata.gz: 75493a10996762aacbaeb5d0ec0a8299565758eb5d1e2d8697353f02c90113b865d31253615864ea5778804830fb6abc17f50c0e5650862e6cc77319f6bc57b7
7
+ data.tar.gz: ea96557baaeed4cf99d9144cabfcadf70e6ac00024bffb16fdcc299da7298c5ec073fb9d91bdfcfbf6cd25bbd71f178bb973c7340720ba429521d94f0d4a769e
data/Rakefile CHANGED
@@ -1,4 +1,6 @@
1
1
  #!/usr/bin/env rake
2
+ # frozen_string_literal: true
3
+
2
4
  begin
3
5
  require 'bundler/setup'
4
6
  rescue LoadError
@@ -19,7 +21,7 @@ RDoc::Task.new(:rdoc) do |rdoc|
19
21
  rdoc.rdoc_files.include('lib/**/*.rb')
20
22
  end
21
23
 
22
- APP_RAKEFILE = File.expand_path("../spec/dummy/Rakefile", __FILE__)
24
+ APP_RAKEFILE = File.expand_path('spec/dummy/Rakefile', __dir__)
23
25
  load 'rails/tasks/engine.rake'
24
26
 
25
27
  Bundler::GemHelper.install_tasks
@@ -27,4 +29,10 @@ Bundler::GemHelper.install_tasks
27
29
  require 'rspec/core/rake_task'
28
30
  RSpec::Core::RakeTask.new
29
31
 
30
- task :default => :spec
32
+ require 'rubocop/rake_task'
33
+
34
+ RuboCop::RakeTask.new do |task|
35
+ task.options = ['-S']
36
+ end
37
+
38
+ task default: %i[spec rubocop]
@@ -6,44 +6,24 @@ require 'switchman/environment'
6
6
  require 'switchman/errors'
7
7
 
8
8
  module Switchman
9
- class Shard < ::ActiveRecord::Base
9
+ class Shard < UnshardedRecord
10
10
  # ten trillion possible ids per shard. yup.
11
11
  IDS_PER_SHARD = 10_000_000_000_000
12
12
 
13
- CATEGORIES =
14
- {
15
- # special cased to mean all other models
16
- :primary => nil,
17
- # special cased to not allow activating a shard other than the default
18
- :unsharded => [Shard]
19
- }
20
- private_constant :CATEGORIES
21
- @connection_specification_name = @shard_category = :unsharded
22
-
23
- if defined?(::ProtectedAttributes)
24
- attr_accessible :default, :name, :database_server
25
- end
26
-
27
13
  # only allow one default
28
- validates_uniqueness_of :default, :if => lambda { |s| s.default? }
14
+ validates_uniqueness_of :default, if: ->(s) { s.default? }
29
15
 
30
16
  after_save :clear_cache
31
17
  after_destroy :clear_cache
32
18
 
33
- after_rollback :on_rollback
34
-
35
19
  scope :primary, -> { where(name: nil).order(:database_server_id, :id).distinct_on(:database_server_id) }
36
20
 
37
21
  class << self
38
- def categories
39
- CATEGORIES.keys
22
+ def sharded_models
23
+ @sharded_models ||= [::ActiveRecord::Base, UnshardedRecord].freeze
40
24
  end
41
25
 
42
- def default(reload_deprecated = false, reload: false, with_fallback: false)
43
- if reload_deprecated
44
- reload = reload_deprecated
45
- ::ActiveSupport::Deprecation.warn("positional reload parameter to Switchman::Shard.default is deprecated; use `reload: true`")
46
- end
26
+ def default(reload: false, with_fallback: false)
47
27
  if !@default || reload
48
28
  # Have to create a dummy object so that several key methods still work
49
29
  # (it's easier to do this in one place here, and just assume that sharding
@@ -54,74 +34,68 @@ module Switchman
54
34
 
55
35
  # if we already have a default shard in place, and the caller wants
56
36
  # to use it as a fallback, use that instead of the dummy instance
57
- if with_fallback && @default
58
- default = @default
59
- end
37
+ default = @default if with_fallback && @default
60
38
 
61
39
  # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
62
40
  @default ||= default
63
41
 
64
42
  # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
65
43
  @default = begin
66
- find_cached("default_shard") { Shard.where(default: true).take } || default
44
+ find_cached('default_shard') { Shard.where(default: true).take } || default
67
45
  rescue
68
46
  default
69
47
  end
70
48
 
71
- # rebuild current shard activations - it might have "another" default shard serialized there
72
- active_shards.replace(active_shards.dup.map do |category, shard|
73
- shard = Shard.lookup((!shard || shard.default?) ? 'default' : shard.id)
74
- [category, shard]
75
- end.to_h)
76
-
77
- activate!(primary: @default) if active_shards.empty?
78
-
79
49
  # make sure this is not erroneously cached
80
- if @default.database_server.instance_variable_defined?(:@primary_shard)
81
- @default.database_server.remove_instance_variable(:@primary_shard)
82
- end
50
+ @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
83
51
 
84
52
  # and finally, check for cached references to the default shard on the existing connection
85
- if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.shard.default?
86
- ::ActiveRecord::Base.connection.shard = @default
53
+ sharded_models.each do |klass|
54
+ klass.connection.shard = @default if klass.connected? && klass.connection.shard.default?
87
55
  end
88
56
  end
89
57
  @default
90
58
  end
91
59
 
92
- def current(category = :primary)
93
- active_shards[category] || Shard.default
60
+ def current(klass = ::ActiveRecord::Base)
61
+ klass ||= ::ActiveRecord::Base
62
+ klass.connection_pool.shard
94
63
  end
95
64
 
96
65
  def activate(shards)
97
- old_shards = activate!(shards)
66
+ activated_classes = activate!(shards)
98
67
  yield
99
68
  ensure
100
- active_shards.merge!(old_shards) if old_shards
69
+ activated_classes.each do |klass|
70
+ klass.connection_pool.shard_stack.pop
71
+ klass.connected_to_stack.pop
72
+ end
101
73
  end
102
74
 
103
75
  def activate!(shards)
104
- old_shards = nil
105
- currently_active_shards = active_shards
106
- shards.each do |category, shard|
107
- next if category == :unsharded
108
- unless currently_active_shards[category] == shard
109
- old_shards ||= {}
110
- old_shards[category] = currently_active_shards[category]
111
- currently_active_shards[category] = shard
112
- end
76
+ activated_classes = []
77
+ shards.each do |klass, shard|
78
+ next if klass == UnshardedRecord
79
+
80
+ next unless klass.current_shard != shard.database_server.id.to_sym ||
81
+ klass.connection_pool.shard != shard
82
+
83
+ activated_classes << klass
84
+ klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass] }
85
+ klass.connection_pool.shard_stack << shard
113
86
  end
114
- old_shards
87
+ activated_classes
115
88
  end
116
89
 
117
90
  def lookup(id)
118
91
  id_i = id.to_i
119
92
  return current if id_i == current.id || id == 'self'
120
93
  return default if id_i == default.id || id.nil? || id == 'default'
94
+
121
95
  id = id_i
122
- raise ArgumentError if id == 0
96
+ raise ArgumentError if id.zero?
123
97
 
124
- unless cached_shards.has_key?(id)
98
+ unless cached_shards.key?(id)
125
99
  cached_shards[id] = Shard.default.activate do
126
100
  find_cached(['shard', id]) { find_by(id: id) }
127
101
  end
@@ -136,59 +110,47 @@ module Switchman
136
110
  # ==== Parameters
137
111
  #
138
112
  # * +shards+ - an array or relation of Shards to iterate over
139
- # * +categories+ - an array of categories to activate
140
- # * +options+ -
141
- # :parallel - true/false to execute in parallel, or a integer of how many
113
+ # * +classes+ - an array of classes to activate
114
+ # parallel: - true/false to execute in parallel, or a integer of how many
142
115
  # sub-processes per database server. Note that parallel
143
116
  # invocation currently uses forking, so should be used sparingly
144
117
  # because errors are not raised, and you cannot get results back
145
- # :max_procs - only run this many parallel processes at a time
146
- # :exception - :ignore, :raise, :defer (wait until the end and raise the first
118
+ # max_procs: - only run this many parallel processes at a time
119
+ # exception: - :ignore, :raise, :defer (wait until the end and raise the first
147
120
  # error), or a proc
148
- def with_each_shard(*args)
149
- raise ArgumentError, "wrong number of arguments (#{args.length} for 0...3)" if args.length > 3
121
+ def with_each_shard(*args, parallel: false, max_procs: nil, exception: :raise, &block)
122
+ raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
150
123
 
151
- unless default.is_a?(Shard)
152
- return Array.wrap(yield)
153
- end
124
+ return Array.wrap(yield) unless default.is_a?(Shard)
154
125
 
155
- options = args.extract_options!
156
126
  if args.length == 1
157
- if Array === args.first && args.first.first.is_a?(Symbol)
158
- categories = args.first
127
+ if Array === args.first && args.first.first.is_a?(Class)
128
+ classes = args.first
159
129
  else
160
130
  scope = args.first
161
131
  end
162
132
  else
163
- scope, categories = args
133
+ scope, classes = args
164
134
  end
165
135
 
166
- parallel = case options[:parallel]
167
- when true
168
- 1
169
- when false, nil
170
- 0
171
- else
172
- options[:parallel]
173
- end
174
- options.delete(:parallel)
136
+ parallel = 1 if parallel == true
137
+ parallel = 0 if parallel == false || parallel.nil?
175
138
 
176
139
  scope ||= Shard.all
177
- if ::ActiveRecord::Relation === scope && scope.order_values.empty?
178
- scope = scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
179
- end
140
+ scope = scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id')) if ::ActiveRecord::Relation === scope && scope.order_values.empty?
180
141
 
181
- if parallel > 0
182
- max_procs = determine_max_procs(options.delete(:max_procs), parallel)
142
+ if parallel.positive?
143
+ max_procs = determine_max_procs(max_procs, parallel)
183
144
  if ::ActiveRecord::Relation === scope
184
145
  # still need a post-uniq, cause the default database server could be NULL or Rails.env in the db
185
146
  database_servers = scope.reorder('database_server_id').select(:database_server_id).distinct.
186
- map(&:database_server).compact.uniq
147
+ map(&:database_server).compact.uniq
187
148
  # nothing to do
188
- return if database_servers.count == 0
149
+ return if database_servers.count.zero?
150
+
189
151
  parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
190
152
 
191
- scopes = Hash[database_servers.map do |server|
153
+ scopes = database_servers.map do |server|
192
154
  server_scope = server.shards.merge(scope)
193
155
  if parallel == 1
194
156
  subscopes = [server_scope]
@@ -196,28 +158,28 @@ module Switchman
196
158
  subscopes = []
197
159
  total = server_scope.count
198
160
  ranges = []
199
- server_scope.find_ids_in_ranges(:batch_size => (total.to_f / parallel).ceil) do |min, max|
161
+ server_scope.find_ids_in_ranges(batch_size: (total.to_f / parallel).ceil) do |min, max|
200
162
  ranges << [min, max]
201
163
  end
202
164
  # create a half-open range on the last one
203
165
  ranges.last[1] = nil
204
166
  ranges.each do |min, max|
205
- subscope = server_scope.where("id>=?", min)
206
- subscope = subscope.where("id<=?", max) if max
167
+ subscope = server_scope.where('id>=?', min)
168
+ subscope = subscope.where('id<=?', max) if max
207
169
  subscopes << subscope
208
170
  end
209
171
  end
210
172
  [server, subscopes]
211
- end]
173
+ end.to_h
212
174
  else
213
175
  scopes = scope.group_by(&:database_server)
214
176
  if parallel > 1
215
177
  parallel = [(max_procs.to_f / scopes.count).ceil, parallel].min if max_procs
216
- scopes = Hash[scopes.map do |(server, shards)|
178
+ scopes = scopes.map do |(server, shards)|
217
179
  [server, shards.in_groups(parallel, false).compact]
218
- end]
180
+ end.to_h
219
181
  else
220
- scopes = Hash[scopes.map { |(server, shards)| [server, [shards]] }]
182
+ scopes = scopes.map { |(server, shards)| [server, [shards]] }.to_h
221
183
  end
222
184
  end
223
185
 
@@ -229,8 +191,8 @@ module Switchman
229
191
  fd_to_name_map = {}
230
192
  errors = []
231
193
 
232
- wait_for_output = lambda do |out_fds, err_fds, fd_to_name_map|
233
- ready, _ = IO.select(out_fds + err_fds)
194
+ wait_for_output = lambda do
195
+ ready, = IO.select(out_fds + err_fds)
234
196
  ready.each do |fd|
235
197
  if fd.eof?
236
198
  fd.close
@@ -245,7 +207,8 @@ module Switchman
245
207
 
246
208
  # only one process; don't bother forking
247
209
  if scopes.length == 1 && parallel == 1
248
- return with_each_shard(scopes.first.last.first, categories, options) { yield }
210
+ return with_each_shard(scopes.first.last.first, classes, exception: exception,
211
+ &block)
249
212
  end
250
213
 
251
214
  # clear connections prior to forking (no more queries will be executed in the parent,
@@ -255,51 +218,47 @@ module Switchman
255
218
 
256
219
  scopes.each do |server, subscopes|
257
220
  subscopes.each_with_index do |subscope, idx|
258
- if subscopes.length > 1
259
- name = "#{server.id} #{idx + 1}"
260
- else
261
- name = server.id
262
- end
221
+ name = if subscopes.length > 1
222
+ "#{server.id} #{idx + 1}"
223
+ else
224
+ server.id
225
+ end
263
226
 
264
227
  exception_pipe = IO.pipe
265
228
  exception_pipes << exception_pipe
266
229
  pid, io_in, io_out, io_err = Open4.pfork4(lambda do
230
+ Switchman.config[:on_fork_proc]&.call
231
+
232
+ # set a pretty name for the process title, up to 128 characters
233
+ # (we don't actually know the limit, depending on how the process
234
+ # was started)
235
+ # first, simplify the binary name by stripping directories,
236
+ # then truncate arguments as necessary
237
+ bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
238
+ max_length = 128 - bin.length - name.length - 3
239
+ args = ARGV.join(' ')
240
+ args = args[0..max_length] if max_length >= 0
241
+ new_title = [bin, args, name].join(' ')
242
+ Process.setproctitle(new_title)
243
+
244
+ with_each_shard(subscope, classes, exception: exception, &block)
245
+ exception_pipe.last.close
246
+ rescue => e
267
247
  begin
268
- Switchman.config[:on_fork_proc]&.call
269
-
270
- # set a pretty name for the process title, up to 128 characters
271
- # (we don't actually know the limit, depending on how the process
272
- # was started)
273
- # first, simplify the binary name by stripping directories,
274
- # then truncate arguments as necessary
275
- bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
276
- max_length = 128 - bin.length - name.length - 3
277
- args = ARGV.join(" ")
278
- if max_length >= 0
279
- args = args[0..max_length]
280
- end
281
- new_title = [bin, args, name].join(" ")
282
- Process.setproctitle(new_title)
283
-
284
- with_each_shard(subscope, categories, options) { yield }
285
- exception_pipe.last.close
286
- rescue => e
287
- begin
288
- dumped = Marshal.dump(e)
289
- rescue
290
- # couldn't dump the exception; create a copy with just
291
- # the message and the backtrace
292
- e2 = e.class.new(e.message)
293
- e2.set_backtrace(e.backtrace)
294
- e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
295
- dumped = Marshal.dump(e2)
296
- end
297
- exception_pipe.last.set_encoding(dumped.encoding)
298
- exception_pipe.last.write(dumped)
299
- exception_pipe.last.flush
300
- exception_pipe.last.close
301
- exit! 1
248
+ dumped = Marshal.dump(e)
249
+ rescue
250
+ # couldn't dump the exception; create a copy with just
251
+ # the message and the backtrace
252
+ e2 = e.class.new(e.message)
253
+ e2.set_backtrace(e.backtrace)
254
+ e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
255
+ dumped = Marshal.dump(e2)
302
256
  end
257
+ exception_pipe.last.set_encoding(dumped.encoding)
258
+ exception_pipe.last.write(dumped)
259
+ exception_pipe.last.flush
260
+ exception_pipe.last.close
261
+ exit! 1
303
262
  end)
304
263
  exception_pipe.last.close
305
264
  pids << pid
@@ -313,7 +272,7 @@ module Switchman
313
272
  while max_procs && pids.count >= max_procs
314
273
  while max_procs && out_fds.count >= max_procs
315
274
  # wait for output if we've hit the max_procs limit
316
- wait_for_output.call(out_fds, err_fds, fd_to_name_map)
275
+ wait_for_output.call
317
276
  end
318
277
  # we've gotten all the output from one fd so wait for its child process to exit
319
278
  found_pid, status = Process.wait2
@@ -323,9 +282,7 @@ module Switchman
323
282
  end
324
283
  end
325
284
 
326
- while out_fds.any? || err_fds.any?
327
- wait_for_output.call(out_fds, err_fds, fd_to_name_map)
328
- end
285
+ wait_for_output.call while out_fds.any? || err_fds.any?
329
286
  pids.each do |pid|
330
287
  _, status = Process.waitpid2(pid)
331
288
  errors << pid_to_name_map[pid] if status.exitstatus != 0
@@ -333,67 +290,51 @@ module Switchman
333
290
 
334
291
  # check for an exception; we only re-raise the first one
335
292
  exception_pipes.each do |exception_pipe|
336
- begin
337
- serialized_exception = exception_pipe.first.read
338
- next if serialized_exception.empty?
339
- exception = Marshal.load(serialized_exception)
340
- raise exception
341
- ensure
342
- exception_pipe.first.close
343
- end
293
+ serialized_exception = exception_pipe.first.read
294
+ next if serialized_exception.empty?
295
+
296
+ ex = Marshal.load(serialized_exception) # rubocop:disable Security/MarshalLoad
297
+ raise ex
298
+ ensure
299
+ exception_pipe.first.close
344
300
  end
345
301
 
346
302
  unless errors.empty?
347
- raise ParallelShardExecError.new("The following subprocesses did not exit cleanly: #{errors.sort.join(", ")}")
303
+ raise ParallelShardExecError,
304
+ "The following subprocesses did not exit cleanly: #{errors.sort.join(', ')}"
348
305
  end
306
+
349
307
  return
350
308
  end
351
309
 
352
- categories ||= []
310
+ classes ||= []
353
311
 
354
312
  previous_shard = nil
355
- close_connections_if_needed = lambda do |shard|
356
- # prune the prior connection unless it happened to be the same
357
- if previous_shard && shard != previous_shard && !previous_shard.database_server.shareable?
358
- previous_shard.activate do
359
- ::GuardRail.activated_environments.each do |env|
360
- ::GuardRail.activate(env) do
361
- if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.open_transactions == 0
362
- ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
363
- end
364
- end
365
- end
366
- end
367
- end
368
- end
369
-
370
313
  result = []
371
- exception = nil
314
+ ex = nil
372
315
  scope.each do |shard|
373
316
  # shard references a database server that isn't configured in this environment
374
317
  next unless shard.database_server
375
- close_connections_if_needed.call(shard)
376
- shard.activate(*categories) do
377
- begin
378
- result.concat Array.wrap(yield)
379
- rescue
380
- case options[:exception]
381
- when :ignore
382
- when :defer
383
- exception ||= $!
384
- when Proc
385
- options[:exception].call
386
- when :raise
387
- raise
388
- else
389
- raise
390
- end
318
+
319
+ shard.activate(*classes) do
320
+ result.concat Array.wrap(yield)
321
+ rescue
322
+ case exception
323
+ when :ignore
324
+ # ignore
325
+ when :defer
326
+ ex ||= $!
327
+ when Proc
328
+ exception.call
329
+ # when :raise
330
+ else
331
+ raise
391
332
  end
392
333
  end
393
334
  previous_shard = shard
394
335
  end
395
- close_connections_if_needed.call(Shard.current)
396
- raise exception if exception
336
+ raise ex if ex
337
+
397
338
  result
398
339
  end
399
340
 
@@ -402,22 +343,22 @@ module Switchman
402
343
  array.each do |object|
403
344
  partition_object = partition_proc ? partition_proc.call(object) : object
404
345
  case partition_object
405
- when Shard
406
- shard = partition_object
407
- when ::ActiveRecord::Base
408
- if partition_object.respond_to?(:associated_shards)
409
- partition_object.associated_shards.each do |a_shard|
410
- shard_arrays[a_shard] ||= []
411
- shard_arrays[a_shard] << object
412
- end
413
- next
414
- else
415
- shard = partition_object.shard
346
+ when Shard
347
+ shard = partition_object
348
+ when ::ActiveRecord::Base
349
+ if partition_object.respond_to?(:associated_shards)
350
+ partition_object.associated_shards.each do |a_shard|
351
+ shard_arrays[a_shard] ||= []
352
+ shard_arrays[a_shard] << object
416
353
  end
417
- when Integer, /^\d+$/, /^(\d+)~(\d+)$/
418
- local_id, shard = Shard.local_id_for(partition_object)
419
- local_id ||= partition_object
420
- object = local_id if !partition_proc
354
+ next
355
+ else
356
+ shard = partition_object.shard
357
+ end
358
+ when Integer, /^\d+$/, /^(\d+)~(\d+)$/
359
+ local_id, shard = Shard.local_id_for(partition_object)
360
+ local_id ||= partition_object
361
+ object = local_id unless partition_proc
421
362
  end
422
363
  shard ||= Shard.current
423
364
  shard_arrays[shard] ||= []
@@ -426,7 +367,7 @@ module Switchman
426
367
  # TODO: use with_each_shard (or vice versa) to get
427
368
  # connection management and parallelism benefits
428
369
  shard_arrays.inject([]) do |results, (shard, objects)|
429
- results.concat shard.activate { Array.wrap(yield objects) }
370
+ results.concat(shard.activate { Array.wrap(yield objects) })
430
371
  end
431
372
  end
432
373
 
@@ -438,7 +379,7 @@ module Switchman
438
379
  # stay as provided. This assumes no consumer
439
380
  # will return a nil value from the block.
440
381
  def signed_id_operation(input_id)
441
- sign = input_id < 0 ? -1 : 1
382
+ sign = input_id.negative? ? -1 : 1
442
383
  output = yield input_id.abs
443
384
  output * sign
444
385
  end
@@ -446,9 +387,10 @@ module Switchman
446
387
  # converts an AR object, integral id, string id, or string short-global-id to a
447
388
  # integral id. nil if it can't be interpreted
448
389
  def integral_id_for(any_id)
449
- if any_id.is_a?(::Arel::Nodes::Casted)
450
- any_id = any_id.val
451
- elsif any_id.is_a?(::Arel::Nodes::BindParam) && ::Rails.version >= "5.2"
390
+ case any_id
391
+ when ::Arel::Nodes::Casted
392
+ any_id = any_id.value
393
+ when ::Arel::Nodes::BindParam
452
394
  any_id = any_id.value.value_before_type_cast
453
395
  end
454
396
 
@@ -459,12 +401,11 @@ module Switchman
459
401
  local_id = $2.to_i
460
402
  signed_id_operation(local_id) do |id|
461
403
  return nil if id > IDS_PER_SHARD
404
+
462
405
  $1.to_i * IDS_PER_SHARD + id
463
406
  end
464
407
  when Integer, /^-?\d+$/
465
408
  any_id.to_i
466
- else
467
- nil
468
409
  end
469
410
  end
470
411
 
@@ -476,11 +417,12 @@ module Switchman
476
417
  def local_id_for(any_id)
477
418
  id = integral_id_for(any_id)
478
419
  return NIL_NIL_ID unless id
420
+
479
421
  return_shard = nil
480
422
  local_id = signed_id_operation(id) do |abs_id|
481
423
  if abs_id < IDS_PER_SHARD
482
424
  abs_id
483
- elsif return_shard = lookup(abs_id / IDS_PER_SHARD)
425
+ elsif (return_shard = lookup(abs_id / IDS_PER_SHARD))
484
426
  abs_id % IDS_PER_SHARD
485
427
  else
486
428
  return NIL_NIL_ID
@@ -497,8 +439,10 @@ module Switchman
497
439
  integral_id = integral_id_for(any_id)
498
440
  local_id, shard = local_id_for(integral_id)
499
441
  return integral_id unless local_id
442
+
500
443
  shard ||= source_shard
501
444
  return local_id if shard == target_shard
445
+
502
446
  shard.global_id_for(local_id)
503
447
  end
504
448
 
@@ -509,6 +453,7 @@ module Switchman
509
453
  local_id, shard = local_id_for(any_id)
510
454
  return any_id unless local_id
511
455
  return local_id unless shard
456
+
512
457
  "#{shard.id}~#{local_id}"
513
458
  end
514
459
 
@@ -517,6 +462,7 @@ module Switchman
517
462
  def global_id_for(any_id, source_shard = nil)
518
463
  id = integral_id_for(any_id)
519
464
  return any_id unless id
465
+
520
466
  signed_id_operation(id) do |abs_id|
521
467
  if abs_id >= IDS_PER_SHARD
522
468
  abs_id
@@ -529,29 +475,59 @@ module Switchman
529
475
 
530
476
  def shard_for(any_id, source_shard = nil)
531
477
  return any_id.shard if any_id.is_a?(::ActiveRecord::Base)
478
+
532
479
  _, shard = local_id_for(any_id)
533
480
  shard || source_shard || Shard.current
534
481
  end
535
482
 
536
483
  # given the provided option, determines whether we need to (and whether
537
484
  # it's possible) to determine a reasonable default.
538
- def determine_max_procs(max_procs_input, parallel_input=2)
485
+ def determine_max_procs(max_procs_input, parallel_input = 2)
539
486
  max_procs = nil
540
487
  if max_procs_input
541
488
  max_procs = max_procs_input.to_i
542
- max_procs = nil if max_procs == 0
489
+ max_procs = nil if max_procs.zero?
543
490
  else
544
491
  return 1 if parallel_input.nil? || parallel_input < 1
492
+
545
493
  cpus = Environment.cpu_count
546
- if cpus && cpus > 0
547
- max_procs = cpus * parallel_input
548
- end
494
+ max_procs = cpus * parallel_input if cpus&.positive?
549
495
  end
550
496
 
551
- return max_procs
497
+ max_procs
552
498
  end
553
499
 
554
500
  private
501
+
502
+ def add_sharded_model(klass)
503
+ @sharded_models = (sharded_models + [klass]).freeze
504
+ initialize_sharding
505
+ end
506
+
507
+ def initialize_sharding
508
+ full_connects_to_hash = DatabaseServer.all.map { |db| [db.id.to_sym, db.connects_to_hash] }.to_h
509
+ sharded_models.each do |klass|
510
+ connects_to_hash = full_connects_to_hash.deep_dup
511
+ if klass == UnshardedRecord
512
+ # no need to mention other databases for the unsharded category
513
+ connects_to_hash = { ::Rails.env.to_sym => DatabaseServer.find(nil).connects_to_hash }
514
+ end
515
+
516
+ # prune things we're already connected to
517
+ if klass.connection_specification_name == klass.name
518
+ connects_to_hash.each do |(db_name, role_hash)|
519
+ role_hash.each_key do |role|
520
+ role_hash.delete(role) if klass.connection_handler.retrieve_connection_pool(
521
+ klass.connection_specification_name, role: role, shard: db_name
522
+ )
523
+ end
524
+ end
525
+ end
526
+
527
+ klass.connects_to shards: connects_to_hash
528
+ end
529
+ end
530
+
555
531
  # in-process caching
556
532
  def cached_shards
557
533
  @cached_shards ||= {}.compare_by_identity
@@ -582,10 +558,6 @@ module Switchman
582
558
  shard = nil unless shard.database_server
583
559
  shard
584
560
  end
585
-
586
- def active_shards
587
- Thread.current[:active_shards] ||= {}.compare_by_identity
588
- end
589
561
  end
590
562
 
591
563
  def name
@@ -599,11 +571,11 @@ module Switchman
599
571
 
600
572
  def name=(name)
601
573
  write_attribute(:name, @name = name)
602
- remove_instance_variable(:@name) if name == nil
574
+ remove_instance_variable(:@name) if name.nil?
603
575
  end
604
576
 
605
577
  def database_server
606
- @database_server ||= DatabaseServer.find(self.database_server_id)
578
+ @database_server ||= DatabaseServer.find(database_server_id)
607
579
  end
608
580
 
609
581
  def database_server=(database_server)
@@ -624,23 +596,21 @@ module Switchman
624
596
  Shard.default
625
597
  end
626
598
 
627
- def activate(*categories)
628
- shards = hashify_categories(categories)
629
- Shard.activate(shards) do
630
- yield
631
- end
599
+ def activate(*classes, &block)
600
+ shards = hashify_classes(classes)
601
+ Shard.activate(shards, &block)
632
602
  end
633
603
 
634
604
  # for use from console ONLY
635
- def activate!(*categories)
636
- shards = hashify_categories(categories)
605
+ def activate!(*classes)
606
+ shards = hashify_classes(classes)
637
607
  Shard.activate!(shards)
638
608
  nil
639
609
  end
640
610
 
641
611
  # custom serialization, since shard is self-referential
642
- def _dump(depth)
643
- self.id.to_s
612
+ def _dump(_depth)
613
+ id.to_s
644
614
  end
645
615
 
646
616
  def self._load(str)
@@ -648,45 +618,45 @@ module Switchman
648
618
  end
649
619
 
650
620
  def drop_database
651
- raise("Cannot drop the database of the default shard") if self.default?
621
+ raise('Cannot drop the database of the default shard') if default?
652
622
  return unless read_attribute(:name)
653
623
 
654
624
  begin
655
- adapter = self.database_server.config[:adapter]
625
+ adapter = database_server.config[:adapter]
656
626
  sharding_config = Switchman.config || {}
657
627
  drop_statement = sharding_config[adapter]&.[](:drop_statement)
658
628
  drop_statement ||= sharding_config[:drop_statement]
659
629
  if drop_statement
660
630
  drop_statement = Array(drop_statement).dup.
661
- map { |statement| statement.gsub('%{name}', self.name) }
631
+ map { |statement| statement.gsub('%{name}', name) }
662
632
  end
663
633
 
664
634
  case adapter
665
- when 'mysql', 'mysql2'
666
- self.activate do
667
- ::GuardRail.activate(:deploy) do
668
- drop_statement ||= "DROP DATABASE #{self.name}"
669
- Array(drop_statement).each do |stmt|
670
- ::ActiveRecord::Base.connection.execute(stmt)
671
- end
635
+ when 'mysql', 'mysql2'
636
+ activate do
637
+ ::GuardRail.activate(:deploy) do
638
+ drop_statement ||= "DROP DATABASE #{name}"
639
+ Array(drop_statement).each do |stmt|
640
+ ::ActiveRecord::Base.connection.execute(stmt)
672
641
  end
673
642
  end
674
- when 'postgresql'
675
- self.activate do
676
- ::GuardRail.activate(:deploy) do
677
- # Shut up, Postgres!
678
- conn = ::ActiveRecord::Base.connection
679
- old_proc = conn.raw_connection.set_notice_processor {}
680
- begin
681
- drop_statement ||= "DROP SCHEMA #{self.name} CASCADE"
682
- Array(drop_statement).each do |stmt|
683
- ::ActiveRecord::Base.connection.execute(stmt)
684
- end
685
- ensure
686
- conn.raw_connection.set_notice_processor(&old_proc) if old_proc
643
+ end
644
+ when 'postgresql'
645
+ activate do
646
+ ::GuardRail.activate(:deploy) do
647
+ # Shut up, Postgres!
648
+ conn = ::ActiveRecord::Base.connection
649
+ old_proc = conn.raw_connection.set_notice_processor {}
650
+ begin
651
+ drop_statement ||= "DROP SCHEMA #{name} CASCADE"
652
+ Array(drop_statement).each do |stmt|
653
+ ::ActiveRecord::Base.connection.execute(stmt)
687
654
  end
655
+ ensure
656
+ conn.raw_connection.set_notice_processor(&old_proc) if old_proc
688
657
  end
689
658
  end
659
+ end
690
660
  end
691
661
  rescue
692
662
  logger.info "Drop failed: #{$!}"
@@ -696,8 +666,9 @@ module Switchman
696
666
  # takes an id local to this shard, and returns a global id
697
667
  def global_id_for(local_id)
698
668
  return nil unless local_id
669
+
699
670
  self.class.signed_id_operation(local_id) do |abs_id|
700
- abs_id + self.id * IDS_PER_SHARD
671
+ abs_id + id * IDS_PER_SHARD
701
672
  end
702
673
  end
703
674
 
@@ -707,7 +678,8 @@ module Switchman
707
678
  end
708
679
 
709
680
  def destroy
710
- raise("Cannot destroy the default shard") if self.default?
681
+ raise('Cannot destroy the default shard') if default?
682
+
711
683
  super
712
684
  end
713
685
 
@@ -716,7 +688,7 @@ module Switchman
716
688
  def clear_cache
717
689
  Shard.default.activate do
718
690
  Switchman.cache.delete(['shard', id].join('/'))
719
- Switchman.cache.delete("default_shard") if default?
691
+ Switchman.cache.delete('default_shard') if default?
720
692
  end
721
693
  self.class.clear_cache
722
694
  end
@@ -725,22 +697,14 @@ module Switchman
725
697
  database_server.shard_name(self)
726
698
  end
727
699
 
728
- def on_rollback
729
- # make sure all connection pool proxies are referencing valid pools
730
- ::ActiveRecord::Base.connection_handler.connection_pools.each do |pool|
731
- next unless pool.is_a?(ConnectionPoolProxy)
732
-
733
- pool.remove_shard!(self)
734
- end
735
- end
736
-
737
- def hashify_categories(categories)
738
- if categories.empty?
739
- { :primary => self }
700
+ def hashify_classes(classes)
701
+ if classes.empty?
702
+ { ::ActiveRecord::Base => self }
740
703
  else
741
- categories.inject({}) { |h, category| h[category] = self; h }
704
+ classes.each_with_object({}) do |klass, h|
705
+ h[klass] = self
706
+ end
742
707
  end
743
708
  end
744
-
745
709
  end
746
710
  end