From 7dd0c1586fa760fbdf8a3195d67596bd6b26b10d Mon Sep 17 00:00:00 2001 From: Cory Johns Date: Thu, 15 Jun 2017 11:55:24 -0400 Subject: [PATCH] Configurable and larger max message size (#146) The default max message / frame size for the websocket library is 1MB. However, the initial AllWatcher response for large models can exceed that. Increased the default to 4MB and made it configurable. Fixes #136 --- juju/client/connection.py | 36 +++++++++++++++++++++++++----------- juju/controller.py | 17 +++++++++++------ juju/model.py | 14 ++++++++++---- tox.ini | 1 + 4 files changed, 47 insertions(+), 21 deletions(-) diff --git a/juju/client/connection.py b/juju/client/connection.py index 6851707..6f2f2a2 100644 --- a/juju/client/connection.py +++ b/juju/client/connection.py @@ -111,9 +111,14 @@ class Connection: Note: Any connection method or constructor can accept an optional `loop` argument to override the default event loop from `asyncio.get_event_loop`. """ + + DEFAULT_FRAME_SIZE = 'default_frame_size' + MAX_FRAME_SIZE = 2**22 + "Maximum size for a single frame. Defaults to 4MB." + def __init__( self, endpoint, uuid, username, password, cacert=None, - macaroons=None, loop=None): + macaroons=None, loop=None, max_frame_size=DEFAULT_FRAME_SIZE): self.endpoint = endpoint self.uuid = uuid if macaroons: @@ -133,6 +138,9 @@ class Connection: self.facades = {} self.messages = IdQueue(loop=self.loop) self.monitor = Monitor(connection=self) + if max_frame_size is self.DEFAULT_FRAME_SIZE: + max_frame_size = self.MAX_FRAME_SIZE + self.max_frame_size = max_frame_size @property def is_open(self): @@ -153,6 +161,7 @@ class Connection: kw = dict() kw['ssl'] = self._get_ssl(self.cacert) kw['loop'] = self.loop + kw['max_size'] = self.max_frame_size self.addr = url self.ws = await websockets.connect(url, **kw) self.loop.create_task(self.receiver()) @@ -321,6 +330,7 @@ class Connection: self.cacert, self.macaroons, self.loop, + self.max_frame_size, ) async def controller(self): @@ -372,7 +382,7 @@ class Connection: @classmethod async def connect( cls, endpoint, uuid, username, password, cacert=None, - macaroons=None, loop=None): + macaroons=None, loop=None, max_frame_size=None): """Connect to the websocket. If uuid is None, the connection will be to the controller. Otherwise it @@ -380,7 +390,7 @@ class Connection: """ client = cls(endpoint, uuid, username, password, cacert, macaroons, - loop) + loop, max_frame_size) endpoints = [(endpoint, cacert)] while endpoints: _endpoint, _cacert = endpoints.pop(0) @@ -402,7 +412,7 @@ class Connection: return client @classmethod - async def connect_current(cls, loop=None): + async def connect_current(cls, loop=None, max_frame_size=None): """Connect to the currently active model. """ @@ -415,10 +425,10 @@ class Connection: model_name = jujudata.current_model() return await cls.connect_model( - '{}:{}'.format(controller_name, model_name), loop) + '{}:{}'.format(controller_name, model_name), loop, max_frame_size) @classmethod - async def connect_current_controller(cls, loop=None): + async def connect_current_controller(cls, loop=None, max_frame_size=None): """Connect to the currently active controller. """ @@ -427,10 +437,12 @@ class Connection: if not controller_name: raise JujuConnectionError('No current controller') - return await cls.connect_controller(controller_name, loop) + return await cls.connect_controller(controller_name, loop, + max_frame_size) @classmethod - async def connect_controller(cls, controller_name, loop=None): + async def connect_controller(cls, controller_name, loop=None, + max_frame_size=None): """Connect to a controller by name. """ @@ -444,10 +456,11 @@ class Connection: macaroons = get_macaroons(controller_name) if not password else None return await cls.connect( - endpoint, None, username, password, cacert, macaroons, loop) + endpoint, None, username, password, cacert, macaroons, loop, + max_frame_size) @classmethod - async def connect_model(cls, model, loop=None): + async def connect_model(cls, model, loop=None, max_frame_size=None): """Connect to a model by name. :param str model: [:] @@ -478,7 +491,8 @@ class Connection: macaroons = get_macaroons(controller_name) if not password else None return await cls.connect( - endpoint, model_uuid, username, password, cacert, macaroons, loop) + endpoint, model_uuid, username, password, cacert, macaroons, loop, + max_frame_size) def build_facades(self, facades): self.facades.clear() diff --git a/juju/controller.py b/juju/controller.py index e279a74..9b452c7 100644 --- a/juju/controller.py +++ b/juju/controller.py @@ -11,7 +11,8 @@ log = logging.getLogger(__name__) class Controller(object): - def __init__(self, loop=None): + def __init__(self, loop=None, + max_frame_size=connection.Connection.DEFAULT_FRAME_SIZE): """Instantiate a new Controller. One of the connect_* methods will need to be called before this @@ -21,6 +22,7 @@ class Controller(object): """ self.loop = loop or asyncio.get_event_loop() + self.max_frame_size = None self.connection = None self.controller_name = None @@ -30,21 +32,24 @@ class Controller(object): """ self.connection = await connection.Connection.connect( - endpoint, None, username, password, cacert, macaroons) + endpoint, None, username, password, cacert, macaroons, + max_frame_size=self.max_frame_size) async def connect_current(self): """Connect to the current Juju controller. """ self.connection = ( - await connection.Connection.connect_current_controller()) + await connection.Connection.connect_current_controller( + max_frame_size=self.max_frame_size)) async def connect_controller(self, controller_name): """Connect to a Juju controller by name. """ self.connection = ( - await connection.Connection.connect_controller(controller_name)) + await connection.Connection.connect_controller( + controller_name, max_frame_size=self.max_frame_size)) self.controller_name = controller_name async def disconnect(self): @@ -238,7 +243,6 @@ class Controller(object): self.connection) return await controller_facade.AllModels() - def get_payloads(self, *patterns): """Return list of known payloads. @@ -297,7 +301,8 @@ class Controller(object): client_facade = client.UserManagerFacade.from_connection( self.connection) user = tag.user(username) - return await client_facade.UserInfo([client.Entity(user)], include_disabled) + return await client_facade.UserInfo([client.Entity(user)], + include_disabled) async def grant(self, username, acl='login'): """Set access level of the given user on the controller diff --git a/juju/model.py b/juju/model.py index 4db711b..61905c9 100644 --- a/juju/model.py +++ b/juju/model.py @@ -382,13 +382,17 @@ class Model(object): """ The main API for interacting with a Juju model. """ - def __init__(self, loop=None): + def __init__(self, loop=None, + max_frame_size=connection.Connection.DEFAULT_FRAME_SIZE): """Instantiate a new connected Model. :param loop: an asyncio event loop + :param max_frame_size: See + `juju.client.connection.Connection.MAX_FRAME_SIZE` """ self.loop = loop or asyncio.get_event_loop() + self.max_frame_size = max_frame_size self.connection = None self.observers = weakref.WeakValueDictionary() self.state = ModelState(self) @@ -416,6 +420,8 @@ class Model(object): """ if 'loop' not in kw: kw['loop'] = self.loop + if 'max_frame_size' not in kw: + kw['max_frame_size'] = self.max_frame_size self.connection = await connection.Connection.connect(*args, **kw) await self._after_connect() @@ -424,7 +430,7 @@ class Model(object): """ self.connection = await connection.Connection.connect_current( - self.loop) + self.loop, max_frame_size=self.max_frame_size) await self._after_connect() async def connect_model(self, model_name): @@ -433,8 +439,8 @@ class Model(object): :param model_name: Format [controller:][user/]model """ - self.connection = await connection.Connection.connect_model(model_name, - self.loop) + self.connection = await connection.Connection.connect_model( + model_name, self.loop, self.max_frame_size) await self._after_connect() async def _after_connect(self): diff --git a/tox.ini b/tox.ini index 5668cdf..789bbeb 100644 --- a/tox.ini +++ b/tox.ini @@ -17,6 +17,7 @@ deps = pytest-xdist mock asynctest + ipdb [testenv:py35] # default tox env excludes integration tests -- 2.25.1