add ChannelMonitor to track unclosed channels with weakrefs, not __del__

This commit is contained in:
Brian Warner 2015-10-06 16:31:41 -07:00
parent fc30fa6cd4
commit e0cc32af9d
3 changed files with 22 additions and 15 deletions

View File

@ -10,6 +10,7 @@ from .. import __version__
from .. import codes from .. import codes
from ..errors import ServerError, Timeout, WrongPasswordError, UsageError from ..errors import ServerError, Timeout, WrongPasswordError, UsageError
from ..util.hkdf import HKDF from ..util.hkdf import HKDF
from ..channel_monitor import monitor
SECOND = 1 SECOND = 1
MINUTE = 60*SECOND MINUTE = 60*SECOND
@ -141,7 +142,6 @@ class Wormhole:
self.verifier = None self.verifier = None
self._sent_data = False self._sent_data = False
self._got_data = False self._got_data = False
self._closed = False
def handle_welcome(self, welcome): def handle_welcome(self, welcome):
if ("motd" in welcome and if ("motd" in welcome and
@ -194,6 +194,7 @@ class Wormhole:
self.code = code self.code = code
channelid = int(mo.group(1)) channelid = int(mo.group(1))
self.channel = self._channel_manager.connect(channelid) self.channel = self._channel_manager.connect(channelid)
monitor.add(self.channel)
def _start(self): def _start(self):
# allocate the rest now too, so it can be serialized # allocate the rest now too, so it can be serialized
@ -263,9 +264,5 @@ class Wormhole:
raise WrongPasswordError raise WrongPasswordError
def close(self): def close(self):
monitor.close(self.channel)
self.channel.deallocate() self.channel.deallocate()
self._closed = True
def __del__(self):
if not self._closed:
print("Error: a Wormhole instance was not closed", file=sys.stderr)

View File

@ -0,0 +1,16 @@
from __future__ import print_function
import sys
from weakref import ref
class ChannelMonitor:
def __init__(self):
self._open_channels = set()
def add(self, w):
wr = ref(w, self._lost)
self._open_channels.add(wr)
def _lost(self, wr):
print("Error: a Wormhole instance was not closed", file=sys.stderr)
def close(self, w):
self._open_channels.discard(ref(w))
monitor = ChannelMonitor() # singleton

View File

@ -15,6 +15,7 @@ from .. import __version__
from .. import codes from .. import codes
from ..errors import ServerError, WrongPasswordError, UsageError from ..errors import ServerError, WrongPasswordError, UsageError
from ..util.hkdf import HKDF from ..util.hkdf import HKDF
from ..channel_monitor import monitor
@implementer(IBodyProducer) @implementer(IBodyProducer)
class DataProducer: class DataProducer:
@ -152,7 +153,6 @@ class Wormhole:
self._started_get_code = False self._started_get_code = False
self._sent_data = False self._sent_data = False
self._got_data = False self._got_data = False
self._closed = False
def _set_side(self, side): def _set_side(self, side):
self._side = side self._side = side
@ -209,6 +209,7 @@ class Wormhole:
self.code = code self.code = code
channelid = int(mo.group(1)) channelid = int(mo.group(1))
self.channel = self._channel_manager.connect(channelid) self.channel = self._channel_manager.connect(channelid)
monitor.add(self.channel)
def _start(self): def _start(self):
# allocate the rest now too, so it can be serialized # allocate the rest now too, so it can be serialized
@ -325,13 +326,6 @@ class Wormhole:
return d return d
def close(self, res=None): def close(self, res=None):
monitor.close(self.channel)
d = self.channel.deallocate() d = self.channel.deallocate()
def _closed(_):
self._closed = True
return res
d.addCallback(_closed)
return d return d
def __del__(self):
if not self._closed:
print("Error: a Wormhole instance was not closed", file=sys.stderr)