647 lines
		
	
	
		
			22 KiB
		
	
	
	
		
			Python
		
	
	
			
		
		
	
	
			647 lines
		
	
	
		
			22 KiB
		
	
	
	
		
			Python
		
	
	
| # -*- coding: utf-8 -*-
 | |
| # Copyright 2017 Vector Creations Ltd
 | |
| #
 | |
| # Licensed under the Apache License, Version 2.0 (the "License");
 | |
| # you may not use this file except in compliance with the License.
 | |
| # You may obtain a copy of the License at
 | |
| #
 | |
| #     http://www.apache.org/licenses/LICENSE-2.0
 | |
| #
 | |
| # Unless required by applicable law or agreed to in writing, software
 | |
| # distributed under the License is distributed on an "AS IS" BASIS,
 | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | |
| # See the License for the specific language governing permissions and
 | |
| # limitations under the License.
 | |
| """This module contains the implementation of both the client and server
 | |
| protocols.
 | |
| 
 | |
| The basic structure of the protocol is line based, where the initial word of
 | |
| each line specifies the command. The rest of the line is parsed based on the
 | |
| command. For example, the `RDATA` command is defined as::
 | |
| 
 | |
|     RDATA <stream_name> <token> <row_json>
 | |
| 
 | |
| (Note that `<row_json>` may contains spaces, but cannot contain newlines.)
 | |
| 
 | |
| Blank lines are ignored.
 | |
| 
 | |
| # Example
 | |
| 
 | |
| An example iteraction is shown below. Each line is prefixed with '>' or '<' to
 | |
| indicate which side is sending, these are *not* included on the wire::
 | |
| 
 | |
|     * connection established *
 | |
|     > SERVER localhost:8823
 | |
|     > PING 1490197665618
 | |
|     < NAME synapse.app.appservice
 | |
|     < PING 1490197665618
 | |
|     < REPLICATE events 1
 | |
|     < REPLICATE backfill 1
 | |
|     < REPLICATE caches 1
 | |
|     > POSITION events 1
 | |
|     > POSITION backfill 1
 | |
|     > POSITION caches 1
 | |
|     > RDATA caches 2 ["get_user_by_id",["@01register-user:localhost:8823"],1490197670513]
 | |
|     > RDATA events 14 ["$149019767112vOHxz:localhost:8823",
 | |
|         "!AFDCvgApUmpdfVjIXm:localhost:8823","m.room.guest_access","",null]
 | |
|     < PING 1490197675618
 | |
|     > ERROR server stopping
 | |
|     * connection closed by server *
 | |
| """
 | |
| 
 | |
| from twisted.internet import defer
 | |
| from twisted.protocols.basic import LineOnlyReceiver
 | |
| from twisted.python.failure import Failure
 | |
| 
 | |
| from commands import (
 | |
|     COMMAND_MAP, VALID_CLIENT_COMMANDS, VALID_SERVER_COMMANDS,
 | |
|     ErrorCommand, ServerCommand, RdataCommand, PositionCommand, PingCommand,
 | |
|     NameCommand, ReplicateCommand, UserSyncCommand, SyncCommand,
 | |
| )
 | |
| from streams import STREAMS_MAP
 | |
| 
 | |
| from synapse.util.stringutils import random_string
 | |
| from synapse.metrics.metric import CounterMetric
 | |
| 
 | |
| import logging
 | |
| import synapse.metrics
 | |
| import struct
 | |
| import fcntl
 | |
| 
 | |
| 
 | |
| metrics = synapse.metrics.get_metrics_for(__name__)
 | |
| 
 | |
| connection_close_counter = metrics.register_counter(
 | |
|     "close_reason", labels=["reason_type"],
 | |
| )
 | |
| 
 | |
| 
 | |
| # A list of all connected protocols. This allows us to send metrics about the
 | |
| # connections.
 | |
| connected_connections = []
 | |
| 
 | |
| 
 | |
| logger = logging.getLogger(__name__)
 | |
| 
 | |
| 
 | |
| PING_TIME = 5000
 | |
| PING_TIMEOUT_MULTIPLIER = 5
 | |
| PING_TIMEOUT_MS = PING_TIME * PING_TIMEOUT_MULTIPLIER
 | |
| 
 | |
| 
 | |
| class ConnectionStates(object):
 | |
|     CONNECTING = "connecting"
 | |
|     ESTABLISHED = "established"
 | |
|     PAUSED = "paused"
 | |
|     CLOSED = "closed"
 | |
| 
 | |
| 
 | |
| class BaseReplicationStreamProtocol(LineOnlyReceiver):
 | |
|     """Base replication protocol shared between client and server.
 | |
| 
 | |
|     Reads lines (ignoring blank ones) and parses them into command classes,
 | |
|     asserting that they are valid for the given direction, i.e. server commands
 | |
|     are only sent by the server.
 | |
| 
 | |
|     On receiving a new command it calls `on_<COMMAND_NAME>` with the parsed
 | |
|     command.
 | |
| 
 | |
|     It also sends `PING` periodically, and correctly times out remote connections
 | |
|     (if they send a `PING` command)
 | |
|     """
 | |
|     delimiter = b'\n'
 | |
| 
 | |
|     VALID_INBOUND_COMMANDS = []  # Valid commands we expect to receive
 | |
|     VALID_OUTBOUND_COMMANDS = []  # Valid commans we can send
 | |
| 
 | |
|     max_line_buffer = 10000
 | |
| 
 | |
|     def __init__(self, clock):
 | |
|         self.clock = clock
 | |
| 
 | |
|         self.last_received_command = self.clock.time_msec()
 | |
|         self.last_sent_command = 0
 | |
|         self.time_we_closed = None  # When we requested the connection be closed
 | |
| 
 | |
|         self.received_ping = False  # Have we reecived a ping from the other side
 | |
| 
 | |
|         self.state = ConnectionStates.CONNECTING
 | |
| 
 | |
|         self.name = "anon"  # The name sent by a client.
 | |
|         self.conn_id = random_string(5)  # To dedupe in case of name clashes.
 | |
| 
 | |
|         # List of pending commands to send once we've established the connection
 | |
|         self.pending_commands = []
 | |
| 
 | |
|         # The LoopingCall for sending pings.
 | |
|         self._send_ping_loop = None
 | |
| 
 | |
|         self.inbound_commands_counter = CounterMetric(
 | |
|             "inbound_commands", labels=["command"],
 | |
|         )
 | |
|         self.outbound_commands_counter = CounterMetric(
 | |
|             "outbound_commands", labels=["command"],
 | |
|         )
 | |
| 
 | |
|     def connectionMade(self):
 | |
|         logger.info("[%s] Connection established", self.id())
 | |
| 
 | |
|         self.state = ConnectionStates.ESTABLISHED
 | |
| 
 | |
|         connected_connections.append(self)  # Register connection for metrics
 | |
| 
 | |
|         self.transport.registerProducer(self, True)  # For the *Producing callbacks
 | |
| 
 | |
|         self._send_pending_commands()
 | |
| 
 | |
|         # Starts sending pings
 | |
|         self._send_ping_loop = self.clock.looping_call(self.send_ping, 5000)
 | |
| 
 | |
|         # Always send the initial PING so that the other side knows that they
 | |
|         # can time us out.
 | |
|         self.send_command(PingCommand(self.clock.time_msec()))
 | |
| 
 | |
|     def send_ping(self):
 | |
|         """Periodically sends a ping and checks if we should close the connection
 | |
|         due to the other side timing out.
 | |
|         """
 | |
|         now = self.clock.time_msec()
 | |
| 
 | |
|         if self.time_we_closed:
 | |
|             if now - self.time_we_closed > PING_TIMEOUT_MS:
 | |
|                 logger.info(
 | |
|                     "[%s] Failed to close connection gracefully, aborting", self.id()
 | |
|                 )
 | |
|                 self.transport.abortConnection()
 | |
|         else:
 | |
|             if now - self.last_sent_command >= PING_TIME:
 | |
|                 self.send_command(PingCommand(now))
 | |
| 
 | |
|             if self.received_ping and now - self.last_received_command > PING_TIMEOUT_MS:
 | |
|                 logger.info(
 | |
|                     "[%s] Connection hasn't received command in %r ms. Closing.",
 | |
|                     self.id(), now - self.last_received_command
 | |
|                 )
 | |
|                 self.send_error("ping timeout")
 | |
| 
 | |
|     def lineReceived(self, line):
 | |
|         """Called when we've received a line
 | |
|         """
 | |
|         if line.strip() == "":
 | |
|             # Ignore blank lines
 | |
|             return
 | |
| 
 | |
|         line = line.decode("utf-8")
 | |
|         cmd_name, rest_of_line = line.split(" ", 1)
 | |
| 
 | |
|         if cmd_name not in self.VALID_INBOUND_COMMANDS:
 | |
|             logger.error("[%s] invalid command %s", self.id(), cmd_name)
 | |
|             self.send_error("invalid command: %s", cmd_name)
 | |
|             return
 | |
| 
 | |
|         self.last_received_command = self.clock.time_msec()
 | |
| 
 | |
|         self.inbound_commands_counter.inc(cmd_name)
 | |
| 
 | |
|         cmd_cls = COMMAND_MAP[cmd_name]
 | |
|         try:
 | |
|             cmd = cmd_cls.from_line(rest_of_line)
 | |
|         except Exception as e:
 | |
|             logger.exception(
 | |
|                 "[%s] failed to parse line %r: %r", self.id(), cmd_name, rest_of_line
 | |
|             )
 | |
|             self.send_error(
 | |
|                 "failed to parse line for  %r: %r (%r):" % (cmd_name, e, rest_of_line)
 | |
|             )
 | |
|             return
 | |
| 
 | |
|         # Now lets try and call on_<CMD_NAME> function
 | |
|         try:
 | |
|             getattr(self, "on_%s" % (cmd_name,))(cmd)
 | |
|         except Exception:
 | |
|             logger.exception("[%s] Failed to handle line: %r", self.id(), line)
 | |
| 
 | |
|     def close(self):
 | |
|         logger.warn("[%s] Closing connection", self.id())
 | |
|         self.time_we_closed = self.clock.time_msec()
 | |
|         self.transport.loseConnection()
 | |
|         self.on_connection_closed()
 | |
| 
 | |
|     def send_error(self, error_string, *args):
 | |
|         """Send an error to remote and close the connection.
 | |
|         """
 | |
|         self.send_command(ErrorCommand(error_string % args))
 | |
|         self.close()
 | |
| 
 | |
|     def send_command(self, cmd, do_buffer=True):
 | |
|         """Send a command if connection has been established.
 | |
| 
 | |
|         Args:
 | |
|             cmd (Command)
 | |
|             do_buffer (bool): Whether to buffer the message or always attempt
 | |
|                 to send the command. This is mostly used to send an error
 | |
|                 message if we're about to close the connection due our buffers
 | |
|                 becoming full.
 | |
|         """
 | |
|         if self.state == ConnectionStates.CLOSED:
 | |
|             logger.debug("[%s] Not sending, connection closed", self.id())
 | |
|             return
 | |
| 
 | |
|         if do_buffer and self.state != ConnectionStates.ESTABLISHED:
 | |
|             self._queue_command(cmd)
 | |
|             return
 | |
| 
 | |
|         self.outbound_commands_counter.inc(cmd.NAME)
 | |
| 
 | |
|         string = "%s %s" % (cmd.NAME, cmd.to_line(),)
 | |
|         if "\n" in string:
 | |
|             raise Exception("Unexpected newline in command: %r", string)
 | |
| 
 | |
|         self.sendLine(string.encode("utf-8"))
 | |
| 
 | |
|         self.last_sent_command = self.clock.time_msec()
 | |
| 
 | |
|     def _queue_command(self, cmd):
 | |
|         """Queue the command until the connection is ready to write to again.
 | |
|         """
 | |
|         logger.debug("[%s] Queing as conn %r, cmd: %r", self.id(), self.state, cmd)
 | |
|         self.pending_commands.append(cmd)
 | |
| 
 | |
|         if len(self.pending_commands) > self.max_line_buffer:
 | |
|             # The other side is failing to keep up and out buffers are becoming
 | |
|             # full, so lets close the connection.
 | |
|             # XXX: should we squawk more loudly?
 | |
|             logger.error("[%s] Remote failed to keep up", self.id())
 | |
|             self.send_command(ErrorCommand("Failed to keep up"), do_buffer=False)
 | |
|             self.close()
 | |
| 
 | |
|     def _send_pending_commands(self):
 | |
|         """Send any queued commandes
 | |
|         """
 | |
|         pending = self.pending_commands
 | |
|         self.pending_commands = []
 | |
|         for cmd in pending:
 | |
|             self.send_command(cmd)
 | |
| 
 | |
|     def on_PING(self, line):
 | |
|         self.received_ping = True
 | |
| 
 | |
|     def on_ERROR(self, cmd):
 | |
|         logger.error("[%s] Remote reported error: %r", self.id(), cmd.data)
 | |
| 
 | |
|     def pauseProducing(self):
 | |
|         """This is called when both the kernel send buffer and the twisted
 | |
|         tcp connection send buffers have become full.
 | |
| 
 | |
|         We don't actually have any control over those sizes, so we buffer some
 | |
|         commands ourselves before knifing the connection due to the remote
 | |
|         failing to keep up.
 | |
|         """
 | |
|         logger.info("[%s] Pause producing", self.id())
 | |
|         self.state = ConnectionStates.PAUSED
 | |
| 
 | |
|     def resumeProducing(self):
 | |
|         """The remote has caught up after we started buffering!
 | |
|         """
 | |
|         logger.info("[%s] Resume producing", self.id())
 | |
|         self.state = ConnectionStates.ESTABLISHED
 | |
|         self._send_pending_commands()
 | |
| 
 | |
|     def stopProducing(self):
 | |
|         """We're never going to send any more data (normally because either
 | |
|         we or the remote has closed the connection)
 | |
|         """
 | |
|         logger.info("[%s] Stop producing", self.id())
 | |
|         self.on_connection_closed()
 | |
| 
 | |
|     def connectionLost(self, reason):
 | |
|         logger.info("[%s] Replication connection closed: %r", self.id(), reason)
 | |
|         if isinstance(reason, Failure):
 | |
|             connection_close_counter.inc(reason.type.__name__)
 | |
|         else:
 | |
|             connection_close_counter.inc(reason.__class__.__name__)
 | |
| 
 | |
|         try:
 | |
|             # Remove us from list of connections to be monitored
 | |
|             connected_connections.remove(self)
 | |
|         except ValueError:
 | |
|             pass
 | |
| 
 | |
|         # Stop the looping call sending pings.
 | |
|         if self._send_ping_loop and self._send_ping_loop.running:
 | |
|             self._send_ping_loop.stop()
 | |
| 
 | |
|         self.on_connection_closed()
 | |
| 
 | |
|     def on_connection_closed(self):
 | |
|         logger.info("[%s] Connection was closed", self.id())
 | |
| 
 | |
|         self.state = ConnectionStates.CLOSED
 | |
|         self.pending_commands = []
 | |
| 
 | |
|         if self.transport:
 | |
|             self.transport.unregisterProducer()
 | |
| 
 | |
|     def __str__(self):
 | |
|         return "ReplicationConnection<name=%s,conn_id=%s,addr=%s>" % (
 | |
|             self.name, self.conn_id, self.addr,
 | |
|         )
 | |
| 
 | |
|     def id(self):
 | |
|         return "%s-%s" % (self.name, self.conn_id)
 | |
| 
 | |
| 
 | |
| class ServerReplicationStreamProtocol(BaseReplicationStreamProtocol):
 | |
|     VALID_INBOUND_COMMANDS = VALID_CLIENT_COMMANDS
 | |
|     VALID_OUTBOUND_COMMANDS = VALID_SERVER_COMMANDS
 | |
| 
 | |
|     def __init__(self, server_name, clock, streamer, addr):
 | |
|         BaseReplicationStreamProtocol.__init__(self, clock)  # Old style class
 | |
| 
 | |
|         self.server_name = server_name
 | |
|         self.streamer = streamer
 | |
|         self.addr = addr
 | |
| 
 | |
|         # The streams the client has subscribed to and is up to date with
 | |
|         self.replication_streams = set()
 | |
| 
 | |
|         # The streams the client is currently subscribing to.
 | |
|         self.connecting_streams = set()
 | |
| 
 | |
|         # Map from stream name to list of updates to send once we've finished
 | |
|         # subscribing the client to the stream.
 | |
|         self.pending_rdata = {}
 | |
| 
 | |
|     def connectionMade(self):
 | |
|         self.send_command(ServerCommand(self.server_name))
 | |
|         BaseReplicationStreamProtocol.connectionMade(self)
 | |
|         self.streamer.new_connection(self)
 | |
| 
 | |
|     def on_NAME(self, cmd):
 | |
|         logger.info("[%s] Renamed to %r", self.id(), cmd.data)
 | |
|         self.name = cmd.data
 | |
| 
 | |
|     def on_USER_SYNC(self, cmd):
 | |
|         self.streamer.on_user_sync(
 | |
|             self.conn_id, cmd.user_id, cmd.is_syncing, cmd.last_sync_ms,
 | |
|         )
 | |
| 
 | |
|     def on_REPLICATE(self, cmd):
 | |
|         stream_name = cmd.stream_name
 | |
|         token = cmd.token
 | |
| 
 | |
|         if stream_name == "ALL":
 | |
|             # Subscribe to all streams we're publishing to.
 | |
|             for stream in self.streamer.streams_by_name.iterkeys():
 | |
|                 self.subscribe_to_stream(stream, token)
 | |
|         else:
 | |
|             self.subscribe_to_stream(stream_name, token)
 | |
| 
 | |
|     def on_FEDERATION_ACK(self, cmd):
 | |
|         self.streamer.federation_ack(cmd.token)
 | |
| 
 | |
|     def on_REMOVE_PUSHER(self, cmd):
 | |
|         self.streamer.on_remove_pusher(cmd.app_id, cmd.push_key, cmd.user_id)
 | |
| 
 | |
|     def on_INVALIDATE_CACHE(self, cmd):
 | |
|         self.streamer.on_invalidate_cache(cmd.cache_func, cmd.keys)
 | |
| 
 | |
|     def on_USER_IP(self, cmd):
 | |
|         self.streamer.on_user_ip(
 | |
|             cmd.user_id, cmd.access_token, cmd.ip, cmd.user_agent, cmd.device_id,
 | |
|             cmd.last_seen,
 | |
|         )
 | |
| 
 | |
|     @defer.inlineCallbacks
 | |
|     def subscribe_to_stream(self, stream_name, token):
 | |
|         """Subscribe the remote to a streams.
 | |
| 
 | |
|         This invloves checking if they've missed anything and sending those
 | |
|         updates down if they have. During that time new updates for the stream
 | |
|         are queued and sent once we've sent down any missed updates.
 | |
|         """
 | |
|         self.replication_streams.discard(stream_name)
 | |
|         self.connecting_streams.add(stream_name)
 | |
| 
 | |
|         try:
 | |
|             # Get missing updates
 | |
|             updates, current_token = yield self.streamer.get_stream_updates(
 | |
|                 stream_name, token,
 | |
|             )
 | |
| 
 | |
|             # Send all the missing updates
 | |
|             for update in updates:
 | |
|                 token, row = update[0], update[1]
 | |
|                 self.send_command(RdataCommand(stream_name, token, row))
 | |
| 
 | |
|             # We send a POSITION command to ensure that they have an up to
 | |
|             # date token (especially useful if we didn't send any updates
 | |
|             # above)
 | |
|             self.send_command(PositionCommand(stream_name, current_token))
 | |
| 
 | |
|             # Now we can send any updates that came in while we were subscribing
 | |
|             pending_rdata = self.pending_rdata.pop(stream_name, [])
 | |
|             for token, update in pending_rdata:
 | |
|                 # Only send updates newer than the current token
 | |
|                 if token > current_token:
 | |
|                     self.send_command(RdataCommand(stream_name, token, update))
 | |
| 
 | |
|             # They're now fully subscribed
 | |
|             self.replication_streams.add(stream_name)
 | |
|         except Exception as e:
 | |
|             logger.exception("[%s] Failed to handle REPLICATE command", self.id())
 | |
|             self.send_error("failed to handle replicate: %r", e)
 | |
|         finally:
 | |
|             self.connecting_streams.discard(stream_name)
 | |
| 
 | |
|     def stream_update(self, stream_name, token, data):
 | |
|         """Called when a new update is available to stream to clients.
 | |
| 
 | |
|         We need to check if the client is interested in the stream or not
 | |
|         """
 | |
|         if stream_name in self.replication_streams:
 | |
|             # The client is subscribed to the stream
 | |
|             self.send_command(RdataCommand(stream_name, token, data))
 | |
|         elif stream_name in self.connecting_streams:
 | |
|             # The client is being subscribed to the stream
 | |
|             logger.debug("[%s] Queuing RDATA %r %r", self.id(), stream_name, token)
 | |
|             self.pending_rdata.setdefault(stream_name, []).append((token, data))
 | |
|         else:
 | |
|             # The client isn't subscribed
 | |
|             logger.debug("[%s] Dropping RDATA %r %r", self.id(), stream_name, token)
 | |
| 
 | |
|     def send_sync(self, data):
 | |
|         self.send_command(SyncCommand(data))
 | |
| 
 | |
|     def on_connection_closed(self):
 | |
|         BaseReplicationStreamProtocol.on_connection_closed(self)
 | |
|         self.streamer.lost_connection(self)
 | |
| 
 | |
| 
 | |
| class ClientReplicationStreamProtocol(BaseReplicationStreamProtocol):
 | |
|     VALID_INBOUND_COMMANDS = VALID_SERVER_COMMANDS
 | |
|     VALID_OUTBOUND_COMMANDS = VALID_CLIENT_COMMANDS
 | |
| 
 | |
|     def __init__(self, client_name, server_name, clock, handler):
 | |
|         BaseReplicationStreamProtocol.__init__(self, clock)
 | |
| 
 | |
|         self.client_name = client_name
 | |
|         self.server_name = server_name
 | |
|         self.handler = handler
 | |
| 
 | |
|         # Map of stream to batched updates. See RdataCommand for info on how
 | |
|         # batching works.
 | |
|         self.pending_batches = {}
 | |
| 
 | |
|     def connectionMade(self):
 | |
|         self.send_command(NameCommand(self.client_name))
 | |
|         BaseReplicationStreamProtocol.connectionMade(self)
 | |
| 
 | |
|         # Once we've connected subscribe to the necessary streams
 | |
|         for stream_name, token in self.handler.get_streams_to_replicate().iteritems():
 | |
|             self.replicate(stream_name, token)
 | |
| 
 | |
|         # Tell the server if we have any users currently syncing (should only
 | |
|         # happen on synchrotrons)
 | |
|         currently_syncing = self.handler.get_currently_syncing_users()
 | |
|         now = self.clock.time_msec()
 | |
|         for user_id in currently_syncing:
 | |
|             self.send_command(UserSyncCommand(user_id, True, now))
 | |
| 
 | |
|         # We've now finished connecting to so inform the client handler
 | |
|         self.handler.update_connection(self)
 | |
| 
 | |
|     def on_SERVER(self, cmd):
 | |
|         if cmd.data != self.server_name:
 | |
|             logger.error("[%s] Connected to wrong remote: %r", self.id(), cmd.data)
 | |
|             self.send_error("Wrong remote")
 | |
| 
 | |
|     def on_RDATA(self, cmd):
 | |
|         try:
 | |
|             row = STREAMS_MAP[cmd.stream_name].ROW_TYPE(*cmd.row)
 | |
|         except Exception:
 | |
|             logger.exception(
 | |
|                 "[%s] Failed to parse RDATA: %r %r",
 | |
|                 self.id(), cmd.stream_name, cmd.row
 | |
|             )
 | |
|             raise
 | |
| 
 | |
|         if cmd.token is None:
 | |
|             # I.e. this is part of a batch of updates for this stream. Batch
 | |
|             # until we get an update for the stream with a non None token
 | |
|             self.pending_batches.setdefault(cmd.stream_name, []).append(row)
 | |
|         else:
 | |
|             # Check if this is the last of a batch of updates
 | |
|             rows = self.pending_batches.pop(cmd.stream_name, [])
 | |
|             rows.append(row)
 | |
| 
 | |
|             self.handler.on_rdata(cmd.stream_name, cmd.token, rows)
 | |
| 
 | |
|     def on_POSITION(self, cmd):
 | |
|         self.handler.on_position(cmd.stream_name, cmd.token)
 | |
| 
 | |
|     def on_SYNC(self, cmd):
 | |
|         self.handler.on_sync(cmd.data)
 | |
| 
 | |
|     def replicate(self, stream_name, token):
 | |
|         """Send the subscription request to the server
 | |
|         """
 | |
|         if stream_name not in STREAMS_MAP:
 | |
|             raise Exception("Invalid stream name %r" % (stream_name,))
 | |
| 
 | |
|         logger.info(
 | |
|             "[%s] Subscribing to replication stream: %r from %r",
 | |
|             self.id(), stream_name, token
 | |
|         )
 | |
| 
 | |
|         self.send_command(ReplicateCommand(stream_name, token))
 | |
| 
 | |
|     def on_connection_closed(self):
 | |
|         BaseReplicationStreamProtocol.on_connection_closed(self)
 | |
|         self.handler.update_connection(None)
 | |
| 
 | |
| 
 | |
| # The following simply registers metrics for the replication connections
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "pending_commands",
 | |
|     lambda: {
 | |
|         (p.name, p.conn_id): len(p.pending_commands)
 | |
|         for p in connected_connections
 | |
|     },
 | |
|     labels=["name", "conn_id"],
 | |
| )
 | |
| 
 | |
| 
 | |
| def transport_buffer_size(protocol):
 | |
|     if protocol.transport:
 | |
|         size = len(protocol.transport.dataBuffer) + protocol.transport._tempDataLen
 | |
|         return size
 | |
|     return 0
 | |
| 
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "transport_send_buffer",
 | |
|     lambda: {
 | |
|         (p.name, p.conn_id): transport_buffer_size(p)
 | |
|         for p in connected_connections
 | |
|     },
 | |
|     labels=["name", "conn_id"],
 | |
| )
 | |
| 
 | |
| 
 | |
| def transport_kernel_read_buffer_size(protocol, read=True):
 | |
|     SIOCINQ = 0x541B
 | |
|     SIOCOUTQ = 0x5411
 | |
| 
 | |
|     if protocol.transport:
 | |
|         fileno = protocol.transport.getHandle().fileno()
 | |
|         if read:
 | |
|             op = SIOCINQ
 | |
|         else:
 | |
|             op = SIOCOUTQ
 | |
|         size = struct.unpack("I", fcntl.ioctl(fileno, op, '\0\0\0\0'))[0]
 | |
|         return size
 | |
|     return 0
 | |
| 
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "transport_kernel_send_buffer",
 | |
|     lambda: {
 | |
|         (p.name, p.conn_id): transport_kernel_read_buffer_size(p, False)
 | |
|         for p in connected_connections
 | |
|     },
 | |
|     labels=["name", "conn_id"],
 | |
| )
 | |
| 
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "transport_kernel_read_buffer",
 | |
|     lambda: {
 | |
|         (p.name, p.conn_id): transport_kernel_read_buffer_size(p, True)
 | |
|         for p in connected_connections
 | |
|     },
 | |
|     labels=["name", "conn_id"],
 | |
| )
 | |
| 
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "inbound_commands",
 | |
|     lambda: {
 | |
|         (k[0], p.name, p.conn_id): count
 | |
|         for p in connected_connections
 | |
|         for k, count in p.inbound_commands_counter.counts.iteritems()
 | |
|     },
 | |
|     labels=["command", "name", "conn_id"],
 | |
| )
 | |
| 
 | |
| metrics.register_callback(
 | |
|     "outbound_commands",
 | |
|     lambda: {
 | |
|         (k[0], p.name, p.conn_id): count
 | |
|         for p in connected_connections
 | |
|         for k, count in p.outbound_commands_counter.counts.iteritems()
 | |
|     },
 | |
|     labels=["command", "name", "conn_id"],
 | |
| )
 |