implement timeouts on service.py and scrape.py's stuff

app/functions.py

@@ -1,5 +1,7 @@
 from bs4 import BeautifulSoup
 import MySQLdb
+from pebble import ProcessPool
+from concurrent.futures import TimeoutError
 import markovify
 import requests
 from Crypto.PublicKey import RSA
@@ -46,7 +48,9 @@ def generate_output(handle):
 		host = cfg['db_host'],
 		user=cfg['db_user'],
 		passwd=cfg['db_pass'],
-		db=cfg['db_name']
+		db=cfg['db_name'],
+		use_unicode=True,
+		charset="utf8mb4"
 	)
 	# print("Generating post for {}".format(handle))
 	dc = db.cursor(MySQLdb.cursors.DictCursor)
@@ -88,7 +92,7 @@ def generate_output(handle):
 	posts = "\n".join(list(sum(c.fetchall(), ())))
 	if len(posts) == 0:
 		print("{} - No posts to learn from.".format(handle))
-		return
+		return bot, None
 
 	if bot['fake_mentions'] == 'never':
 		# remove all mentions from the training data before the markov model sees it
@@ -99,8 +103,8 @@ def generate_output(handle):
 	post = None
 
 	# even with such a high tries value for markovify, it still sometimes returns none.
-	# so we implement our own tries function as well, and try ten times.
-	while post is None and tries < 10:
+	# so we implement our own tries function as well, and try five times.
+	while post is None and tries < 5:
 		post = model.make_short_sentence(bot['length'], tries = 1000)
 		tries += 1
 
@@ -138,6 +142,10 @@ def make_post(args):
 
 	bot, post = generate_output(handle)
 
+	# post will be None if there's no posts for the bot to learn from.
+	# in such a case, we should just exit without doing anything.
+	if post == None: return
+
 	client = Mastodon(
 		client_id = bot['client_id'],
 		client_secret = bot['client_secret'],
@@ -149,7 +157,9 @@ def make_post(args):
 		host = cfg['db_host'],
 		user=cfg['db_user'],
 		passwd=cfg['db_pass'],
-		db=cfg['db_name']
+		db=cfg['db_name'],
+		use_unicode=True,
+		charset="utf8mb4"
 	)
 	c = db.cursor()
 
@@ -173,7 +183,7 @@ def make_post(args):
 		# this needs to be dealt with properly later on, but for now, we'll just disable the bot
 		c.execute("UPDATE bots SET enabled = FALSE WHERE handle = %s", (handle,))
 	except:
-		print("Failed to create post for {}".format(handle))
+		print("Failed to submit post for {}".format(handle))
 
 	if id == None:
 		# this wasn't a reply, it was a regular post, so update the last post date
@@ -181,12 +191,31 @@ def make_post(args):
 		db.commit()
 	c.close()
 
+def do_in_pool(function, data, timeout=30, silent=False):
+	with ProcessPool(max_workers=cfg['service_threads']) as p:
+		index = 0
+		future = p.map(function, data)
+		iterator = future.result()
+
+		while True:
+			try:
+				result = next(iterator)
+			except StopIteration:
+				# all threads are done
+				break
+			except TimeoutError as error:
+				if not silent: print("Timed out on {}.".format(data[index]))
+			finally:
+				index += 1
+
 def get_key():
 	db = MySQLdb.connect(
 		host = cfg['db_host'],
 		user=cfg['db_user'],
 		passwd=cfg['db_pass'],
-		db=cfg['db_name']
+		db=cfg['db_name'],
+		use_unicode=True,
+		charset="utf8mb4"
 	)
 
 	dc = db.cursor(MySQLdb.cursors.DictCursor)

app/scrape.py

@@ -2,7 +2,6 @@
 
 import MySQLdb
 import requests
-from multiprocessing import Pool
 import json, re
 import functions
 
@@ -131,7 +130,7 @@ cursor.execute("SELECT `handle`, `outbox` FROM `fedi_accounts` ORDER BY RAND()")
 accounts = cursor.fetchall()
 cursor.close()
 db.close()
-with Pool(cfg['service_threads']) as p:
-	p.map(scrape_posts, accounts)
+
+functions.do_in_pool(scrape_posts, accounts, timeout=60)
 
 print("Done!")

app/service.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
+import json
+
 import MySQLdb
 from mastodon import Mastodon
-from multiprocessing import Pool
 import requests
-import json
+
 import functions
 
 cfg = json.load(open('config.json'))
@@ -73,8 +74,7 @@ cursor.execute("SELECT handle FROM bots WHERE enabled = TRUE AND TIMESTAMPDIFF(M
 # cursor.execute("SELECT handle FROM bots WHERE enabled = TRUE")
 bots = cursor.fetchall()
 
-with Pool(cfg['service_threads']) as p:
-	p.map(functions.make_post, bots)
+functions.do_in_pool(functions.make_post, bots, 15)
 
 print("Updating cached icons")
 dc = db.cursor(MySQLdb.cursors.DictCursor)
@@ -86,7 +86,6 @@ ON bots.credentials_id = credentials.id
 WHERE TIMESTAMPDIFF(HOUR, icon_update_time, CURRENT_TIMESTAMP()) > 2""")
 bots = dc.fetchall()
 
-with Pool(cfg['service_threads']) as p:
-	p.map(update_icon, bots)
+functions.do_in_pool(update_icon, bots)
 
 db.commit()

requirements.txt

@@ -9,3 +9,4 @@ requests==2.23.0
 http-ece==1.1.0
 pycryptodome==3.9.7
 cryptography==2.9.2
+pebble==4.5.3