Add pool for IMAP connections

1 files changed, 129 insertions, 0 deletions

diff --git a/src/acit/imap_pool.py b/src/acit/imap_pool.py
new file mode 100644
index 0000000..7a7017b
--- /dev/null
+++ b/src/acit/imap_pool.py
@@ -0,0 +1,129 @@
+
+from typing import List
+
+from imaplib import IMAP4
+
+import imap_tools
+from imap_tools.errors import MailboxLoginError, MailboxLogoutError
+from imap_tools import MailBox, MailMessage
+
+import socket
+
+from threading import Event, RLock
+
+class PooledMailBox(MailBox):
+	def __init__(self,pool,*args,**kwargs):
+		MailBox.__init__(*args,**kwargs)
+		self._mailbox_pool=pool
+
+	def __enter__(self):
+		return self
+
+	def __exit__(self, exc_type, exc_value, traceback):
+		self.logout()
+	
+	def logout(self):
+		self._mailbox_pool.release(self)
+
+
+
+class MailBoxPool():
+	"A pool of mailboxes."
+	def __init__(self, username:str, password:str, initial_folder:str='INBOX', # for MailBox.login()
+			connection_n=4, # for ourselves
+			host='', port=993, timeout=None, # for MailBox.__init__()
+			keyfile=None, certfile=None, ssl_context=None # for MailBox.__init()
+			):
+		self.n=connection_n
+		self.username=username
+		self.password=password
+		self.initial_folder=initial_folder
+		self.host=host
+		self.port=port
+		self.timeout=timeout
+		self.keyfile=keyfile
+		self.certfile=certfile
+		self.ssl_context=ssl_context
+
+		self.errors=[]
+
+		self.lock=RLock()
+		self.boxreturned=Event()
+		self.pool:List[MailBox]=[]
+		self.taken:List[MailBox]=[]
+
+		self.ensure_all_connections()
+	
+	def is_alive(self,mb:MailBox):
+		"Internal function. Uses the provided mailbox to send a command to the server, checking whether it's still connected"
+		try:
+			mb.client.noop()
+		except (
+			TimeoutError, ConnectionError, # generic python errors
+			IMAP4.error, IMAP4.abort, # IMAP4 errors
+			):
+			return False
+		return True
+
+	def ensure_all_connections(self):
+		"Internal function. Refils the pool with connections. Note this also removes disconnected boxes that are currently in use."
+		with self.lock:
+			toremove=[]
+			for mb in self.pool+self.taken:
+				if not self.is_alive(mb):
+					toremove.append(mb)
+
+			for mb in toremove:
+				try:
+					mb.logout() # logout if needed
+				except:
+					pass
+				if mb in self.pool: self.pool.remove(mb)
+				if mb in self.taken: self.taken.remove(mb)
+
+			boxesleft=len(self.pool)+len(self.taken)
+
+			for i in range(self.n -  boxesleft):
+				try:
+					mb=self.new_box()
+					self.boxreturned.set()
+				except Exception as e:
+					self.errors.append(e)
+	
+	def new_box(self):
+		"Interal function. Generates a new box. Note: this does NOT add it to the pool."
+		mb=MailBox(host=self.host,port=self.port,timeout=self.timeout,
+			certfile=self.certfile,keyfile=self.keyfile,ssl_context=self.ssl_context)
+
+		mb.login(username=self.username,password=self.password,initial_folder=self.initial_folder)
+		return mb
+
+	__public_=["get_pool_size","get_box","release", "errors"]
+	
+	def get_pool_size(self):
+		"Returns the total pool size, free and taken connections."
+		return len(self.pool)+len(self.taken)
+	
+	def get_box(self):
+		"Gets a new mailbox from the pool"
+		while True:
+			with self.lock:
+				self.ensure_all_connections()
+				if len(self.pool)>0:
+					self.boxreturned.clear()
+					mb=self.pool.pop(0)
+					self.taken.append(mb)
+					return mb
+
+			self.boxreturned.wait()
+	
+	def release(self,mb:PooledMailBox):
+		"Returns a mailbox back to the pool. Please use a context manager instead of manually releasing."
+		with self.lock:
+			if mb in self.taken:
+				self.taken.remove(mb)
+			self.pool.append(mb)
+			self.boxreturned.set()
+
+
+