Skip to content

Commit b75fd31

Browse files
miguelgrinbergmiguelgrinberg
committed
Fix admin instrumentation to support disconnect reasons (Fixes #1423)
1 parent 8964dab commit b75fd31

File tree

2 files changed

+68
-96
lines changed

2 files changed

+68
-96
lines changed

‎src/socketio/admin.py‎

Lines changed: 34 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,9 @@ def instrument(self):
7777
# track socket connection times
7878
self.sio.manager._timestamps = {}
7979

80-
# report socket.io connections
81-
self.sio.manager.__connect = self.sio.manager.connect
82-
self.sio.manager.connect = self._connect
83-
84-
# report socket.io disconnection
85-
self.sio.manager.__disconnect = self.sio.manager.disconnect
86-
self.sio.manager.disconnect = self._disconnect
80+
# report socket.io connections, disconnections and received events
81+
self.sio.__trigger_event = self.sio._trigger_event
82+
self.sio._trigger_event = self._trigger_event
8783

8884
# report join rooms
8985
self.sio.manager.__basic_enter_room = \
@@ -99,10 +95,6 @@ def instrument(self):
9995
self.sio.manager.__emit = self.sio.manager.emit
10096
self.sio.manager.emit = self._emit
10197

102-
# report receive events
103-
self.sio.__handle_event_internal = self.sio._handle_event_internal
104-
self.sio._handle_event_internal = self._handle_event_internal
105-
10698
# report engine.io connections
10799
self.sio.eio.on('connect', self._handle_eio_connect)
108100
self.sio.eio.on('disconnect', self._handle_eio_disconnect)
@@ -128,14 +120,12 @@ def instrument(self):
128120

129121
def uninstrument(self): # pragma: no cover
130122
if self.mode == 'development':
131-
self.sio.manager.connect = self.sio.manager.__connect
132-
self.sio.manager.disconnect = self.sio.manager.__disconnect
123+
self.sio._trigger_event = self.sio.__trigger_event
133124
self.sio.manager.basic_enter_room = \
134125
self.sio.manager.__basic_enter_room
135126
self.sio.manager.basic_leave_room = \
136127
self.sio.manager.__basic_leave_room
137128
self.sio.manager.emit = self.sio.manager.__emit
138-
self.sio._handle_event_internal = self.sio.__handle_event_internal
139129
self.sio.eio._ok = self.sio.eio.__ok
140130

141131
from engineio.socket import Socket
@@ -205,26 +195,34 @@ def shutdown(self):
205195
self.stop_stats_event.set()
206196
self.stats_task.join()
207197

208-
def _connect(self, eio_sid, namespace):
209-
sid = self.sio.manager.__connect(eio_sid, namespace)
198+
def _trigger_event(self, event, namespace, *args):
210199
t = time.time()
211-
self.sio.manager._timestamps[sid] = t
212-
serialized_socket = self.serialize_socket(sid, namespace, eio_sid)
213-
self.sio.emit('socket_connected', (
214-
serialized_socket,
215-
datetime.utcfromtimestamp(t).isoformat() + 'Z',
216-
), namespace=self.admin_namespace)
217-
return sid
218-
219-
def _disconnect(self, sid, namespace, **kwargs):
220-
del self.sio.manager._timestamps[sid]
221-
self.sio.emit('socket_disconnected', (
222-
namespace,
223-
sid,
224-
'N/A',
225-
datetime.utcnow().isoformat() + 'Z',
226-
), namespace=self.admin_namespace)
227-
return self.sio.manager.__disconnect(sid, namespace, **kwargs)
200+
sid = args[0]
201+
if event == 'connect':
202+
eio_sid = self.sio.manager.eio_sid_from_sid(sid, namespace)
203+
self.sio.manager._timestamps[sid] = t
204+
serialized_socket = self.serialize_socket(sid, namespace, eio_sid)
205+
self.sio.emit('socket_connected', (
206+
serialized_socket,
207+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
208+
), namespace=self.admin_namespace)
209+
elif event == 'disconnect':
210+
del self.sio.manager._timestamps[sid]
211+
reason = args[1]
212+
self.sio.emit('socket_disconnected', (
213+
namespace,
214+
sid,
215+
reason,
216+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
217+
), namespace=self.admin_namespace)
218+
else:
219+
self.sio.emit('event_received', (
220+
namespace,
221+
sid,
222+
(event, *args[1:]),
223+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
224+
), namespace=self.admin_namespace)
225+
return self.sio.__trigger_event(event, namespace, *args)
228226

229227
def _check_for_upgrade(self, eio_sid, sid, namespace): # pragma: no cover
230228
for _ in range(5):
@@ -269,7 +267,7 @@ def _emit(self, event, data, namespace, room=None, skip_sid=None,
269267
**kwargs)
270268
if namespace != self.admin_namespace:
271269
event_data = [event] + list(data) if isinstance(data, tuple) \
272-
else [data]
270+
else [event, data]
273271
if not isinstance(skip_sid, list): # pragma: no branch
274272
skip_sid = [skip_sid]
275273
for sid, _ in self.sio.manager.get_participants(namespace, room):
@@ -282,18 +280,6 @@ def _emit(self, event, data, namespace, room=None, skip_sid=None,
282280
), namespace=self.admin_namespace)
283281
return ret
284282

285-
def _handle_event_internal(self, server, sid, eio_sid, data, namespace,
286-
id):
287-
ret = self.sio.__handle_event_internal(server, sid, eio_sid, data,
288-
namespace, id)
289-
self.sio.emit('event_received', (
290-
namespace,
291-
sid,
292-
data,
293-
datetime.utcnow().isoformat() + 'Z',
294-
), namespace=self.admin_namespace)
295-
return ret
296-
297283
def _handle_eio_connect(self, eio_sid, environ):
298284
if self.stop_stats_event is None:
299285
self.stop_stats_event = self.sio.eio.create_event()
@@ -303,9 +289,9 @@ def _handle_eio_connect(self, eio_sid, environ):
303289
self.event_buffer.push('rawConnection')
304290
return self.sio._handle_eio_connect(eio_sid, environ)
305291

306-
def _handle_eio_disconnect(self, eio_sid):
292+
def _handle_eio_disconnect(self, eio_sid, reason):
307293
self.event_buffer.push('rawDisconnection')
308-
return self.sio._handle_eio_disconnect(eio_sid)
294+
return self.sio._handle_eio_disconnect(eio_sid, reason)
309295

310296
def _eio_http_response(self, packets=None, headers=None, jsonp_index=None):
311297
ret = self.sio.eio.__ok(packets=packets, headers=headers,

‎src/socketio/async_admin.py‎

Lines changed: 34 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,9 @@ def instrument(self):
5858
# track socket connection times
5959
self.sio.manager._timestamps = {}
6060

61-
# report socket.io connections
62-
self.sio.manager.__connect = self.sio.manager.connect
63-
self.sio.manager.connect = self._connect
64-
65-
# report socket.io disconnection
66-
self.sio.manager.__disconnect = self.sio.manager.disconnect
67-
self.sio.manager.disconnect = self._disconnect
61+
# report socket.io connections, disconnections and received events
62+
self.sio.__trigger_event = self.sio._trigger_event
63+
self.sio._trigger_event = self._trigger_event
6864

6965
# report join rooms
7066
self.sio.manager.__basic_enter_room = \
@@ -80,10 +76,6 @@ def instrument(self):
8076
self.sio.manager.__emit = self.sio.manager.emit
8177
self.sio.manager.emit = self._emit
8278

83-
# report receive events
84-
self.sio.__handle_event_internal = self.sio._handle_event_internal
85-
self.sio._handle_event_internal = self._handle_event_internal
86-
8779
# report engine.io connections
8880
self.sio.eio.on('connect', self._handle_eio_connect)
8981
self.sio.eio.on('disconnect', self._handle_eio_disconnect)
@@ -109,14 +101,12 @@ def instrument(self):
109101

110102
def uninstrument(self): # pragma: no cover
111103
if self.mode == 'development':
112-
self.sio.manager.connect = self.sio.manager.__connect
113-
self.sio.manager.disconnect = self.sio.manager.__disconnect
104+
self.sio._trigger_event = self.sio.__trigger_event
114105
self.sio.manager.basic_enter_room = \
115106
self.sio.manager.__basic_enter_room
116107
self.sio.manager.basic_leave_room = \
117108
self.sio.manager.__basic_leave_room
118109
self.sio.manager.emit = self.sio.manager.__emit
119-
self.sio._handle_event_internal = self.sio.__handle_event_internal
120110
self.sio.eio._ok = self.sio.eio.__ok
121111

122112
from engineio.async_socket import AsyncSocket
@@ -193,26 +183,34 @@ async def shutdown(self):
193183
self.stop_stats_event.set()
194184
await asyncio.gather(self.stats_task)
195185

196-
async def _connect(self, eio_sid, namespace):
197-
sid = await self.sio.manager.__connect(eio_sid, namespace)
186+
async def _trigger_event(self, event, namespace, *args):
198187
t = time.time()
199-
self.sio.manager._timestamps[sid] = t
200-
serialized_socket = self.serialize_socket(sid, namespace, eio_sid)
201-
await self.sio.emit('socket_connected', (
202-
serialized_socket,
203-
datetime.utcfromtimestamp(t).isoformat() + 'Z',
204-
), namespace=self.admin_namespace)
205-
return sid
206-
207-
async def _disconnect(self, sid, namespace, **kwargs):
208-
del self.sio.manager._timestamps[sid]
209-
await self.sio.emit('socket_disconnected', (
210-
namespace,
211-
sid,
212-
'N/A',
213-
datetime.utcnow().isoformat() + 'Z',
214-
), namespace=self.admin_namespace)
215-
return await self.sio.manager.__disconnect(sid, namespace, **kwargs)
188+
sid = args[0]
189+
if event == 'connect':
190+
eio_sid = self.sio.manager.eio_sid_from_sid(sid, namespace)
191+
self.sio.manager._timestamps[sid] = t
192+
serialized_socket = self.serialize_socket(sid, namespace, eio_sid)
193+
await self.sio.emit('socket_connected', (
194+
serialized_socket,
195+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
196+
), namespace=self.admin_namespace)
197+
elif event == 'disconnect':
198+
del self.sio.manager._timestamps[sid]
199+
reason = args[1]
200+
await self.sio.emit('socket_disconnected', (
201+
namespace,
202+
sid,
203+
reason,
204+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
205+
), namespace=self.admin_namespace)
206+
else:
207+
await self.sio.emit('event_received', (
208+
namespace,
209+
sid,
210+
(event, *args[1:]),
211+
datetime.utcfromtimestamp(t).isoformat() + 'Z',
212+
), namespace=self.admin_namespace)
213+
return await self.sio.__trigger_event(event, namespace, *args)
216214

217215
async def _check_for_upgrade(self, eio_sid, sid,
218216
namespace): # pragma: no cover
@@ -258,7 +256,7 @@ async def _emit(self, event, data, namespace, room=None, skip_sid=None,
258256
callback=callback, **kwargs)
259257
if namespace != self.admin_namespace:
260258
event_data = [event] + list(data) if isinstance(data, tuple) \
261-
else [data]
259+
else [event, data]
262260
if not isinstance(skip_sid, list): # pragma: no branch
263261
skip_sid = [skip_sid]
264262
for sid, _ in self.sio.manager.get_participants(namespace, room):
@@ -271,18 +269,6 @@ async def _emit(self, event, data, namespace, room=None, skip_sid=None,
271269
), namespace=self.admin_namespace)
272270
return ret
273271

274-
async def _handle_event_internal(self, server, sid, eio_sid, data,
275-
namespace, id):
276-
ret = await self.sio.__handle_event_internal(server, sid, eio_sid,
277-
data, namespace, id)
278-
await self.sio.emit('event_received', (
279-
namespace,
280-
sid,
281-
data,
282-
datetime.utcnow().isoformat() + 'Z',
283-
), namespace=self.admin_namespace)
284-
return ret
285-
286272
async def _handle_eio_connect(self, eio_sid, environ):
287273
if self.stop_stats_event is None:
288274
self.stop_stats_event = self.sio.eio.create_event()
@@ -292,9 +278,9 @@ async def _handle_eio_connect(self, eio_sid, environ):
292278
self.event_buffer.push('rawConnection')
293279
return await self.sio._handle_eio_connect(eio_sid, environ)
294280

295-
async def _handle_eio_disconnect(self, eio_sid):
281+
async def _handle_eio_disconnect(self, eio_sid, reason):
296282
self.event_buffer.push('rawDisconnection')
297-
return await self.sio._handle_eio_disconnect(eio_sid)
283+
return await self.sio._handle_eio_disconnect(eio_sid, reason)
298284

299285
def _eio_http_response(self, packets=None, headers=None, jsonp_index=None):
300286
ret = self.sio.eio.__ok(packets=packets, headers=headers,

0 commit comments

Comments
 (0)