Rework spawning and task structure

Still not working: See TestBot.py

TestBot.py

@@ -6,41 +6,36 @@ from yaboli.utils import *
 
 #class TestBot(Bot):
 class TestBot(yaboli.Controller):
-	def __init__(self, roomname):
+	def __init__(self, nick):
-		super().__init__(roomname)
+		super().__init__(nick=nick)
-	
-	async def on_snapshot(self, user_id, session_id, version, listing, log, nick=None,
-	                      pm_with_nick=None, pm_with_user_id=None):
-		await self.room.nick("TestBot")
 	
 	async def on_send(self, message):
-		await self.room.send("Hey, a message!", message.message_id)
+		if message.content == "!spawnevil":
+			bot = TestBot("TestSpawn")
+			task, reason = await bot.connect("test")
+			second = await self.room.send("We have " + ("a" if task else "no") + " task. Reason: " + reason, message.message_id)
+			if task:
+				await bot.stop()
+				await self.room.send("Stopped." if task.done() else "Still running (!)", second.message_id)
 		
-	async def on_join(self, session):
+			await self.room.send("All's over now.", message.message_id)
-		if session.nick != "":
-			await self.room.send(f"Hey, a @{mention(session.nick)}!")
-		else:
-			await self.room.send("Hey, a lurker!")
 		
-	async def on_nick(self, session_id, user_id, from_nick, to_nick):
+		elif message.content == "!tree":
-		if from_nick != "" and to_nick != "":
+			messages = [message]
-			if from_nick == to_nick:
+			newmessages = []
-				await self.room.send(f"You didn't even change your nick, @{mention(to_nick)} :(")
+			for i in range(2):
-			else:
+				for m in messages:
-				await self.room.send(f"Bye @{mention(from_nick)}, hi @{mention(to_nick)}")
+					for j in range(2):
-		elif from_nick != "":
+						newm = await self.room.send(f"{m.content}.{j}", m.message_id)
-			await self.room.send(f"Bye @{mention(from_nick)}? This message should never appear...")
+						newmessages.append(newm)
-		elif to_nick != "":
+				messages = newmessages
-			await self.room.send(f"Hey, a @{mention(to_nick)}!")
+				newmessages = []
-		else:
-			await self.room.send("I have no idea how you did that. This message should never appear...")
 
-	async def on_part(self, session):
+async def run_bot():
-		if session.nick != "":
+	bot = TestBot("TestSummoner")
-			await self.room.send(f"Bye, you @{mention(session.nick)}!")
+	task, reason = await bot.connect("test")
-		else:
+	if task:
-			await self.room.send("Bye, you lurker!")
+		await task
 
 if __name__ == "__main__":
-	bot = TestBot("test")
+	asyncio.get_event_loop().run_until_complete(run_bot())
-	asyncio.get_event_loop().run_until_complete(bot.run())

yaboli/__init__.py

@@ -1,3 +1,10 @@
+import logging
+#logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 from .connection import *
 from .room import *
 from .controller import *

yaboli/connection.py

@@ -1,5 +1,6 @@
 import logging
-logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
 import asyncio
 asyncio.get_event_loop().set_debug(True)
 
@@ -17,37 +18,67 @@ class Connection:
 		self.cookie = cookie
 		self.packet_hook = packet_hook
 		
-		self.stopped = False
-		
 		self._ws = None
-		self._pid = 0
+		self._pid = 0 # successive packet ids
+		self._spawned_tasks = set()
 		self._pending_responses = {}
+		#self._stopping = False
+		self._runtask = None
 	
-	async def run(self):
+	async def connect(self, max_tries=10, delay=60):
+		"""
+		success = await connect(max_tries=10, delay=60)
+		
+		Attempt to connect to a room.
+		Returns the task listening for packets, or None if the attempt failed.
+		"""
+		
+		await self.stop()
+		
+		tries_left = max_tries
+		while tries_left > 0:
+			tries_left -= 1
+			try:
 				self._ws = await websockets.connect(self.url, max_size=None)
+			except (websockets.InvalidURI, websockets.InvalidHandshake):
+				self._ws = None
+				if tries_left > 0:
+					await asyncio.sleep(delay)
+			else:
+				self._runtask = asyncio.ensure_future(self._run())
+				return self._runtask
+	
+	async def _run(self):
+		"""
+		Listen for packets and deal with them accordingly.
+		"""
 		
 		try:
 			while True:
-				response = await self._ws.recv()
+				await self._handle_next_message()
-				asyncio.ensure_future(self._handle_json(response))
 		except websockets.ConnectionClosed:
 			pass
 		finally:
-			await self._ws.close() # just to make sure it's closed
+			self._clean_up_futures()
-			self._ws = None
+			self._clean_up_tasks()
-			stopped = True
 			
-			for future in self._pending_responses:
+			await self._ws.close() # just to make sure
-				#future.set_error(ConnectionClosed)
+			self._ws = None
-				future.cancel()
 	
 	async def stop(self):
-		if not self.stopped and self._ws:
+		"""
+		Close websocket connection and wait for running task to stop.
+		"""
+		
+		if self._ws:
 			await self._ws.close()
 		
+		if self._runtask:
+			await self._runtask
+	
 	async def send(self, ptype, data=None, await_response=True):
-		if self.stopped:
+		if not self._ws:
-			raise ConnectionClosed
+			raise asyncio.CancelledError
 		
 		pid = str(self._new_pid())
 		packet = {
@@ -60,7 +91,8 @@ class Connection:
 		if await_response:
 			wait_for = self._wait_for_response(pid)
 		
-		await self._ws.send(json.dumps(packet, separators=(',', ':')))
+		logging.debug(f"Currently used websocket at self._ws: {self._ws}")
+		await self._ws.send(json.dumps(packet, separators=(',', ':'))) # minimum size
 		
 		if await_response:
 			await wait_for
@@ -70,11 +102,32 @@ class Connection:
 		self._pid += 1
 		return self._pid
 	
+	async def _handle_next_message(self):
+		response = await self._ws.recv()
+		task = asyncio.ensure_future(self._handle_json(response))
+		self._track_task(task) # will be cancelled when the connection is closed
+	
+	def _clean_up_futures(self):
+		for pid, future in self._pending_responses.items():
+			logger.debug(f"Cancelling future: {future}")
+			future.cancel()
+		self._pending_responses = {}
+	
+	def _clean_up_tasks(self):
+		for task in self._spawned_tasks:
+			if not task.done():
+				logger.debug(f"Cancelling task: {task}")
+				task.cancel()
+			else:
+				logger.debug(f"Task already done: {task}")
+				logger.debug(f"Exception: {task.exception()}")
+		self._spawned_tasks = set()
+	
 	async def _handle_json(self, text):
 		packet = json.loads(text)
 		
 		# Deal with pending responses
-		pid = packet.get("id")
+		pid = packet.get("id", None)
 		future = self._pending_responses.pop(pid, None)
 		if future:
 			future.set_result(packet)
@@ -82,6 +135,20 @@ class Connection:
 		# Pass packet onto room
 		await self.packet_hook(packet)
 	
+	def _track_task(self, task):
+		self._spawned_tasks.add(task)
+		
+		# only keep running tasks
+		#tasks = set()
+		#for task in self._spawned_tasks:
+			#if not task.done():
+				#logger.debug(f"Keeping task: {task}")
+				#tasks.add(task)
+			#else:
+				#logger.debug(f"Deleting task: {task}")
+		#self._spawned_tasks = tasks
+		#self._spawned_tasks = {task for task in self._spawned_tasks if not task.done()} # TODO: Reenable
+	
 	def _wait_for_response(self, pid):
 		future = asyncio.Future()
 		self._pending_responses[pid] = future

yaboli/controller.py

@@ -1,5 +1,8 @@
+import asyncio
+import logging
 from .room import Room
 
+logger = logging.getLogger(__name__)
 __all__ = ["Controller"]
 
 
@@ -24,50 +27,103 @@ class Controller:
 	
 	"""
 	
-	def __init__(self, roomname, human=False, cookie=None):
+	def __init__(self, nick, human=False, cookie=None):
 		"""
 		roomname  - name of room to connect to
 		human     - whether the human flag should be set on connections
 		cookie    - cookie to use in HTTP request, if any
 		"""
-		
+		self.nick = nick
-		self.roomname = roomname
 		self.human = human
 		self.cookie = cookie
 		
+		self.roomname = "test"
+		self.password = None
+		
 		self.room = None
-		self.running = True
+		self._connect_result = None
 	
-	async def run(self):
+	def _create_room(self, roomname):
-		await self.on_start()
+		return Room(roomname, self, human=self.human, cookie=self.cookie)
 	
-		while self.running:
+	def _set_connect_result(self, result):
-			self.room = Room(self.roomname, self, self.human, self.cookie)
+		logger.debug(f"Attempting to set connect result to {result}")
-			await self.room.run()
+		if self._connect_result and not self._connect_result.done():
+			logger.debug(f"Setting connect result to {result}")
+			self._connect_result.set_result(result)
+	
+	async def connect(self, roomname, password=None, timeout=10):
+		"""
+		task, reason = await connect(roomname, password=None, timeout=10)
+		
+		Connect to a room and authenticate, if necessary.
+		
+		roomname - name of the room to connect to
+		password - password for the room, if needed
+		timeout  - wait this long for a reply from the server
+		
+		Returns:
+		task     - the task running the bot, or None on failure
+		reason   - the reason for failure
+		  "no room"        = could not establish connection, room doesn't exist
+		  "auth option"    = can't authenticate with a password
+		  "no password"    = password needed to connect to room
+		  "wrong password" = password given does not work
+		  "disconnected"   = connection closed before client could access the room
+		  "success"        = no failure
+		"""
+		
+		logger.info(f"Attempting to connect to &{roomname}")
+		
+		# make sure nothing is running any more
+		try:
+			await self.stop()
+		except asyncio.CancelledError:
+			logger.error("Calling connect from the controller itself.")
+			raise
+		
+		self.password = password
+		self.room = self._create_room(roomname)
+		
+		# prepare for if connect() is successful
+		self._connect_result = asyncio.Future()
+		
+		# attempt to connect to the room
+		task = await self.room.connect()
+		if not task:
+			logger.warn(f"Could not connect to &{roomname}.")
 			self.room = None
+			return None, "no room"
 		
-		await self.on_end()
+		# connection succeeded, now we need to know whether we can log in
+		# wait for success/authentication/disconnect
+		# TODO: add a timeout
+		await self._connect_result
+		result = self._connect_result.result()
+		logger.debug(f"&{roomname}._connect_result: {result!r}")
+		
+		# deal with result
+		if result == "success":
+			logger.info(f"Successfully connected to &{roomname}.")
+			return task, result
+		else: # not successful for some reason
+			logger.warn(f"Could not join &{roomname}: {result!r}")
+			await self.stop()
+			return None, result
 	
 	async def stop(self):
-		if self.running:
-			self.running = False
-			
 		if self.room:
+			logger.info(f"&{self.room.roomname}: Stopping")
 			await self.room.stop()
+			logger.debug(f"&{self.room.roomname}: Stopped. Deleting room")
+			self.room = None
 	
-	async def on_start(self):
+	async def set_nick(self, nick):
-		"""
+		if nick != self.nick:
-		The first callback called when the controller is run.
+			_, _, _, to_nick = await self.room.nick(nick)
-		"""
 			
-		pass
+			if to_nick != nick:
-	
+				logger.warn(f"&{self.room.roomname}: Could not set nick to {nick!r}, set to {to_nick!r} instead.")
-	async def on_stop(self):
-		"""
-		The last callback called when the controller is run.
-		"""
-		
-		pass
 	
 	async def on_connected(self):
 		"""
@@ -77,7 +133,7 @@ class Controller:
 		     such as resetting the message history.
 		"""
 		
-		pass
+		self._set_connect_result("success")
 	
 	async def on_disconnected(self):
 		"""
@@ -88,10 +144,18 @@ class Controller:
 		     Need to store information from old room?
 		"""
 		
-		pass
+		logger.debug(f"on_disconnected: self.room is {self.room}")
+		self._set_connect_result("disconnected")
 	
-	async def on_bounce(self, reason=None, auth_options=None, agent_id=None, ip=None):
+	async def on_bounce(self, reason=None, auth_options=[], agent_id=None, ip=None):
-		pass
+		if "passcode" not in auth_options:
+			self._set_connect_result("auth option")
+		elif self.password is None:
+			self._set_connect_result("no password")
+		else:
+			success, reason = await self.room.auth("passcode", passcode=self.password)
+			if not success:
+				self._set_connect_result("wrong password")
 	
 	async def on_disconnect(self, reason):
 		pass
@@ -126,6 +190,7 @@ class Controller:
 		Default implementation, refer to api.euphoria.io
 		"""
 		
+		logger.debug(f"&{self.room.roomname}: Pong!")
 		await self.room.ping_reply(ptime)
 	
 	async def on_pm_initiate(self, from_id, from_nick, from_room, pm_id):
@@ -136,4 +201,5 @@ class Controller:
 	
 	async def on_snapshot(self, user_id, session_id, version, listing, log, nick=None,
 	                      pm_with_nick=None, pm_with_user_id=None):
-		pass
+		if nick != self.nick:
+			await self.room.nick(self.nick)

yaboli/room.py

@@ -1,7 +1,9 @@
 import asyncio
+import logging
 from .connection import *
 from .utils import *
 
+logger = logging.getLogger(__name__)
 __all__ = ["Room"]
 
 
@@ -34,18 +36,41 @@ class Room:
 		self._callbacks = {}
 		self._add_callbacks()
 		
+		self._stopping = False
+		self._runtask = None
+		
 		if human:
 			url = self.HUMAN_FORMAT.format(self.roomname)
 		else:
 			url = self.ROOM_FORMAT.format(self.roomname)
 		self._conn = Connection(url, self._handle_packet, self.cookie)
 	
-	async def run(self):
+	async def connect(self, max_tries=10, delay=60):
-		await self._conn.run()
+		task = await self._conn.connect(max_tries=1)
+		if task:
+			self._runtask = asyncio.ensure_future(self._run(task, max_tries=max_tries, delay=delay))
+			return self._runtask
+	
+	async def _run(self, task, max_tries=10, delay=60):
+		while not self._stopping:
+			await task
+			await self.controller.on_disconnected()
+			
+			task = await self._conn.connect(max_tries=max_tries, delay=delay)
+			if not task:
+				return
+		
+		self.stopping = False
 	
 	async def stop(self):
+		self._stopping = True
 		await self._conn.stop()
 		
+		if self._runtask:
+			await self._runtask
+	
+	
+	
 	# CATEGORY: SESSION COMMANDS
 	
 	async def auth(self, atype, passcode=None):
@@ -270,12 +295,11 @@ class Room:
 			try:
 				await callback(packet)
 			except asyncio.CancelledError as e:
-				# TODO: log error
+				logger.info(f"&{self.roomname}: Callback of type {ptype!r} cancelled.")
-				print("HEHEHEHEY, CANCELLEDERROR", e)
-				pass
 	
 	def _check_for_errors(self, packet):
-		# TODO: log throttled
+		if packet.get("throttled", False):
+			logger.warn(f"&{self.roomname}: Throttled for reason: {packet.get('throttled_reason', 'no reason')!r}")
 		
 		if "error" in packet:
 			raise ResponseError(response.get("error"))