Class: Protobuf::Nats::Server

Inherits:
Object
  • Object
show all
Includes:
Logging, Rpc::Server
Defined in:
lib/protobuf/nats/server.rb

Constant Summary collapse

MILLISECOND = 
1000

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(options) ⇒ Server

Returns a new instance of Server.



17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/protobuf/nats/server.rb', line 17

def initialize(options)
  @options = options
  @processing_requests = true
  @running = true
  @stopped = false
  @pause_mutex = ::Mutex.new

  @nats = @options[:client] || ::Protobuf::Nats::NatsClient.new
  @nats.connect(::Protobuf::Nats.config.connection_options)

  @thread_pool = ::Protobuf::Nats::ThreadPool.new(threads, :max_queue => max_queue_size)

  @subscription_manager = ::Protobuf::Nats::SuperSubscriptionManager.new(@nats) do |request_data, reply_id, subject|
    unless enqueue_request(request_data, reply_id)
      logger.error { "Thread pool is full! Dropping message for subject: #{subject}" }
    end
  end
  @server = options.fetch(:server, ::Socket.gethostname)
end

Instance Attribute Details

#nats ⇒ Object (readonly)

Returns the value of attribute nats.



13
14
15
 
# File 'lib/protobuf/nats/server.rb', line 13

def nats
  @nats
end

#subscription_manager ⇒ Object (readonly)

Returns the value of attribute subscription_manager.



13
14
15
 
# File 'lib/protobuf/nats/server.rb', line 13

def subscription_manager
  @subscription_manager
end

#thread_pool ⇒ Object (readonly)

Returns the value of attribute thread_pool.



13
14
15
 
# File 'lib/protobuf/nats/server.rb', line 13

def thread_pool
  @thread_pool
end

Instance Method Details

#detect_and_handle_a_pause ⇒ Object 



191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
 
# File 'lib/protobuf/nats/server.rb', line 191

def detect_and_handle_a_pause
  @pause_mutex.synchronize do
    case
    # If we are taking requests and detect a pause file, then unsubscribe.
    when @processing_requests && paused?
      @processing_requests = false
      logger.warn("Pausing server!")
      unsubscribe

    # If we were paused and the pause file is no longer present, then subscribe again.
    when !@processing_requests && !paused?
      logger.warn("Resuming server: resubscribing to all services and restarting slow start!")
      @processing_requests = true
      subscribe
    end
  end
end

#do_not_subscribe_to_includes?(subscription_key) ⇒ Boolean

Returns:

  • (Boolean) 


111
112
113
114
115
116
117
118
 
# File 'lib/protobuf/nats/server.rb', line 111

def do_not_subscribe_to_includes?(subscription_key)
  return false unless ::Protobuf::Nats.config.server_subscription_key_do_not_subscribe_to_when_includes_any_of.respond_to?(:any?)
  return false if ::Protobuf::Nats.config.server_subscription_key_do_not_subscribe_to_when_includes_any_of.empty?

  ::Protobuf::Nats.config.server_subscription_key_do_not_subscribe_to_when_includes_any_of.any? do |key|
    subscription_key.include?(key)
  end
end

#enqueue_request(request_data, reply_id) ⇒ Object 



63
64
65
66
67
68
69
70
71
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/protobuf/nats/server.rb', line 63

def enqueue_request(request_data, reply_id)
  ::ActiveSupport::Notifications.instrument "server.message_received.protobuf-nats"

  enqueued_at = ::Time.now
  was_enqueued = thread_pool.push do
    begin
      # Instrument the thread pool time-to-execute duration.
      processed_at = ::Time.now
      ::ActiveSupport::Notifications.instrument("server.thread_pool_execution_delay.protobuf-nats",
                                                (processed_at - enqueued_at) * MILLISECOND)

      # Process request.
      response_data = handle_request(request_data, 'server' => @server)

      # Publish response.
      logger.debug { "Publishing response to #{reply_id}" } if logger.debug?
      nats.publish(reply_id, response_data)
    rescue => error
      logger.debug { "rescued error => #{error}" }  if logger.debug?
      ::Protobuf::Nats.notify_error_callbacks(error)
    ensure
      # Instrument the request duration.
      completed_at = ::Time.now
      ::ActiveSupport::Notifications.instrument("server.request_duration.protobuf-nats",
                                                (completed_at - enqueued_at) * MILLISECOND)
    end
  end

  # Publish an ACK to signal the server has picked up the work.
  begin
    if was_enqueued
      logger.debug { "[reply_id=#{reply_id}] Sending ACK" } if logger.debug?
      nats.publish(reply_id, ::Protobuf::Nats::Messages::ACK)
    else # Drop message if the thread pool is full
      ::ActiveSupport::Notifications.instrument "server.message_dropped.protobuf-nats"
      logger.debug { "[reply_id=#{reply_id}] Sending NACK" } if logger.debug?

      # Let the client know we are not processing the message.
      nats.publish(reply_id, ::Protobuf::Nats::Messages::NACK)
    end
  rescue => e
    logger.error "Failed to send ACK/NACK for #{reply_id}: #{e.message}"
    ::Protobuf::Nats.notify_error_callbacks(e)
  end

  was_enqueued
end

#finish_slow_start ⇒ Object

Slow start subscriptions by adding X rounds of subz every Y seconds, where X is subscriptions_per_rpc_endpoint and Y is slow_start_delay.



166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
 
# File 'lib/protobuf/nats/server.rb', line 166

def finish_slow_start
  logger.info "Slow start has started..."
  completed = 1

  # We have (X - 1) here because we always subscribe at least once.
  (subscriptions_per_rpc_endpoint - 1).times do
    unless @running
      logger.info "Slow start interrupted (server stopping) after #{completed}/#{subscriptions_per_rpc_endpoint} rounds"
      return
    end

    if paused?
      logger.info "Slow start interrupted (server paused) after #{completed}/#{subscriptions_per_rpc_endpoint} rounds"
      return
    end

    completed += 1
    sleep slow_start_delay
    subscribe_to_services_once
    logger.info "Slow start adding another round of subscriptions (#{completed}/#{subscriptions_per_rpc_endpoint})..."
  end

  logger.info "Slow start finished successfully (#{completed}/#{subscriptions_per_rpc_endpoint} rounds completed)."
end

#instrument_thread_pool_sizes ⇒ Object 



37
38
39
40
41
 
# File 'lib/protobuf/nats/server.rb', line 37

def instrument_thread_pool_sizes
  ::ActiveSupport::Notifications.instrument("server.thread_pool_enqueued_size.protobuf-nats", thread_pool.enqueued_size)
  ::ActiveSupport::Notifications.instrument("server.thread_pool_max_size.protobuf-nats", thread_pool.max_size)
  ::ActiveSupport::Notifications.instrument("server.thread_pool_running_size.protobuf-nats", thread_pool.size)
end

#max_queue_size ⇒ Object 



43
44
45
 
# File 'lib/protobuf/nats/server.rb', line 43

def max_queue_size
  ::ENV.fetch("PB_NATS_SERVER_MAX_QUEUE_SIZE", @options[:threads]).to_i
end

#only_subscribe_to_includes?(subscription_key) ⇒ Boolean

Returns:

  • (Boolean) 


120
121
122
123
124
125
126
127
 
# File 'lib/protobuf/nats/server.rb', line 120

def only_subscribe_to_includes?(subscription_key)
  return true unless ::Protobuf::Nats.config.server_subscription_key_only_subscribe_to_when_includes_any_of.respond_to?(:any?)
  return true if ::Protobuf::Nats.config.server_subscription_key_only_subscribe_to_when_includes_any_of.empty?

  ::Protobuf::Nats.config.server_subscription_key_only_subscribe_to_when_includes_any_of.any? do |key|
    subscription_key.include?(key)
  end
end

#pause_file_path ⇒ Object 



129
130
131
 
# File 'lib/protobuf/nats/server.rb', line 129

def pause_file_path
  ::ENV.fetch("PB_NATS_SERVER_PAUSE_FILE_PATH", nil)
end

#paused? ⇒ Boolean

Returns:

  • (Boolean) 


209
210
211
 
# File 'lib/protobuf/nats/server.rb', line 209

def paused?
  !pause_file_path.nil? && ::File.exist?(pause_file_path)
end

#print_subscription_keys ⇒ Object 



133
134
135
136
137
138
139
 
# File 'lib/protobuf/nats/server.rb', line 133

def print_subscription_keys
  logger.info "Creating subscriptions:"

  with_each_subscription_key do |subscription_key|
    logger.info "  - #{subscription_key}"
  end
end

#run ⇒ Object 



213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
 
# File 'lib/protobuf/nats/server.rb', line 213

def run
  nats.on_reconnect do
    logger.warn "Server NATS connection was reconnected"
  end

  nats.on_disconnect do
    logger.warn "Server NATS connection was disconnected"
  end

  nats.on_error do |error|
    ::Protobuf::Nats.notify_error_callbacks(error)
  end

  nats.on_close do
    logger.warn "Server NATS connection was closed"
  end

  print_subscription_keys
  if paused?
    yield if block_given?
  else
    subscribe { yield if block_given? }
  end

  loop do
    break unless @running
    detect_and_handle_a_pause
    instrument_thread_pool_sizes
    sleep 1
  end

  unsubscribe

  logger.info "Shutting down subscription manager..."
  begin
    Timeout.timeout(10) do
      subscription_manager.shutdown(5)
    end
  rescue Timeout::Error
    logger.error "Subscription manager shutdown timed out!"
  rescue => e
    logger.error "Error during subscription manager shutdown: #{e.message}"
  end

  logger.info "Waiting up to 60 seconds for the thread pool to finish shutting down..."
  thread_pool.shutdown
  unless thread_pool.wait_for_termination(60)
    logger.warn "Thread pool did not shut down cleanly within 60 seconds!"
    ::ActiveSupport::Notifications.instrument "server.thread_pool_shutdown_timeout.protobuf-nats"
  end
ensure
  @stopped = true

  begin
    logger.info "Closing NATS connection..."
    @nats.close if @nats
  rescue => e
    logger.warn "Failed to close NATS connection: #{e.message}"
  end
end

#running? ⇒ Boolean

Returns:

  • (Boolean) 


274
275
276
 
# File 'lib/protobuf/nats/server.rb', line 274

def running?
  !@stopped
end

#service_klasses ⇒ Object 



59
60
61
 
# File 'lib/protobuf/nats/server.rb', line 59

def service_klasses
  ::Protobuf::Rpc::Service.implemented_services.map(&:safe_constantize)
end

#slow_start_delay ⇒ Object 



47
48
49
 
# File 'lib/protobuf/nats/server.rb', line 47

def slow_start_delay
  @slow_start_delay ||= ::ENV.fetch("PB_NATS_SERVER_SLOW_START_DELAY", 10).to_i
end

#stop ⇒ Object 



278
279
280
 
# File 'lib/protobuf/nats/server.rb', line 278

def stop
  @running = false
end

#subscribe ⇒ Object 



282
283
284
285
286
 
# File 'lib/protobuf/nats/server.rb', line 282

def subscribe
  subscribe_to_services_once
  yield if block_given?
  finish_slow_start
end

#subscribe_to_services_once ⇒ Object 



141
142
143
144
145
 
# File 'lib/protobuf/nats/server.rb', line 141

def subscribe_to_services_once
  with_each_subscription_key do |subscription_key_and_queue|
    subscription_manager.queue_subscribe(subscription_key_and_queue)
  end
end

#subscriptions_per_rpc_endpoint ⇒ Object 



51
52
53
 
# File 'lib/protobuf/nats/server.rb', line 51

def subscriptions_per_rpc_endpoint
  @subscriptions_per_rpc_endpoint ||= ::ENV.fetch("PB_NATS_SERVER_SUBSCRIPTIONS_PER_RPC_ENDPOINT", 10).to_i
end

#threads ⇒ Object 



55
56
57
 
# File 'lib/protobuf/nats/server.rb', line 55

def threads
  @options[:threads] || 10 # Default to 10 if not provided, consistent with original behavior
end

#unsubscribe ⇒ Object 



288
289
290
291
 
# File 'lib/protobuf/nats/server.rb', line 288

def unsubscribe
  logger.info "Unsubscribing from rpc routes..."
  subscription_manager.unsubscribe_all
end

#with_each_subscription_key ⇒ Object 



147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
 
# File 'lib/protobuf/nats/server.rb', line 147

def with_each_subscription_key
  fail ::ArgumentError unless block_given?

  service_klasses.each do |service_klass|
    service_klass.rpcs.each do |service_method, _|
      # Skip services that are not implemented.
      next unless service_klass.method_defined?(service_method)
      subscription_key = ::Protobuf::Nats.subscription_key(service_klass, service_method)
      next if do_not_subscribe_to_includes?(subscription_key)
      next unless only_subscribe_to_includes?(subscription_key)

      yield subscription_key
    end
  end
end