From 55afdd2905dd1d9f58982691a404c373e768d304 Mon Sep 17 00:00:00 2001
From: Mikhail Galanin <mgalanin@mimecast.com>
Date: Fri, 17 Aug 2018 11:03:57 +0100
Subject: [PATCH] [Minor] Added callback helpers

---
 src/lua/lua_common.h      |  6 ++----
 src/lua/lua_thread_pool.c | 27 ++++++++++++++++++++++-----
 src/lua/lua_thread_pool.h | 28 ++++++++++++++++++++++++++--
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 40d752704..cf171c313 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -419,16 +419,14 @@ gint
 lua_yield_thread (struct thread_entry *thread_entry, gint nresults);
 
 /**
- *
- * @param pool
+ * Resumes suspended by lua_yield_thread () thread
+ * @param task
  * @param thread_entry
  * @param narg
- * @return
  */
 void
 lua_resume_thread (struct rspamd_task *task, struct thread_entry *thread_entry, gint narg);
 
-
 /* Paths defs */
 #define RSPAMD_CONFDIR_INDEX "CONFDIR"
 #define RSPAMD_RUNDIR_INDEX "RUNDIR"
diff --git a/src/lua/lua_thread_pool.c b/src/lua/lua_thread_pool.c
index 865b41fce..979b31f6b 100644
--- a/src/lua/lua_thread_pool.c
+++ b/src/lua/lua_thread_pool.c
@@ -61,7 +61,7 @@ lua_thread_pool_free (struct lua_thread_pool *pool)
 }
 
 struct thread_entry *
-lua_thread_pool_get(struct lua_thread_pool *pool)
+lua_thread_pool_get (struct lua_thread_pool *pool)
 {
 	gpointer cur;
 	struct thread_entry *ent = NULL;
@@ -81,7 +81,7 @@ lua_thread_pool_get(struct lua_thread_pool *pool)
 }
 
 void
-lua_thread_pool_return(struct lua_thread_pool *pool, struct thread_entry *thread_entry)
+lua_thread_pool_return (struct lua_thread_pool *pool, struct thread_entry *thread_entry)
 {
 	g_assert (lua_status (thread_entry->lua_state) == 0); /* we can't return a running/yielded thread into the pool */
 
@@ -99,7 +99,7 @@ lua_thread_pool_return(struct lua_thread_pool *pool, struct thread_entry *thread
 }
 
 void
-lua_thread_pool_terminate_entry(struct lua_thread_pool *pool, struct thread_entry *thread_entry)
+lua_thread_pool_terminate_entry (struct lua_thread_pool *pool, struct thread_entry *thread_entry)
 {
 	struct thread_entry *ent = NULL;
 
@@ -119,13 +119,30 @@ lua_thread_pool_terminate_entry(struct lua_thread_pool *pool, struct thread_entr
 }
 
 struct thread_entry *
-lua_thread_pool_get_running_entry(struct lua_thread_pool *pool)
+lua_thread_pool_get_running_entry (struct lua_thread_pool *pool)
 {
 	return pool->running_entry;
 }
 
 void
-lua_thread_pool_set_running_entry(struct lua_thread_pool *pool, struct thread_entry *thread_entry)
+lua_thread_pool_set_running_entry (struct lua_thread_pool *pool, struct thread_entry *thread_entry)
 {
 	pool->running_entry = thread_entry;
 }
+
+
+void
+lua_thread_pool_prepare_callback (struct lua_thread_pool *pool, struct lua_callback_state *cbs)
+{
+	cbs->thread_pool = pool;
+	cbs->previous_thread = lua_thread_pool_get_running_entry (pool);
+	cbs->my_thread = lua_thread_pool_get (pool);
+	cbs->L = cbs->my_thread->lua_state;
+}
+
+void
+lua_thread_pool_restore_callback (struct lua_callback_state *cbs)
+{
+	lua_thread_pool_return (cbs->thread_pool, cbs->my_thread);
+	lua_thread_pool_set_running_entry (cbs->thread_pool, cbs->previous_thread);
+}
diff --git a/src/lua/lua_thread_pool.h b/src/lua/lua_thread_pool.h
index 64708c488..b72b72e8d 100644
--- a/src/lua/lua_thread_pool.h
+++ b/src/lua/lua_thread_pool.h
@@ -11,6 +11,13 @@ struct thread_entry {
 
 struct thread_pool;
 
+struct lua_callback_state {
+	lua_State *L;
+	struct thread_entry *my_thread;
+	struct thread_entry *previous_thread;
+	struct lua_thread_pool *thread_pool;
+};
+
 /**
  * Allocates new thread pool on state L. Pre-creates number of lua-threads to use later on
  *
@@ -67,7 +74,7 @@ lua_thread_pool_terminate_entry(struct lua_thread_pool *pool, struct thread_entr
  * @return
  */
 struct thread_entry *
-lua_thread_pool_get_running_entry(struct lua_thread_pool *pool);
+lua_thread_pool_get_running_entry (struct lua_thread_pool *pool);
 
 /**
  * Updates currently running thread
@@ -76,7 +83,24 @@ lua_thread_pool_get_running_entry(struct lua_thread_pool *pool);
  * @param thread_entry
  */
 void
-lua_thread_pool_set_running_entry(struct lua_thread_pool *pool, struct thread_entry *thread_entry);
+lua_thread_pool_set_running_entry (struct lua_thread_pool *pool, struct thread_entry *thread_entry);
+
+/**
+ * Prevents yielded thread to be used for callback execution. lua_thread_pool_restore_callback() should be called afterwards.
+ *
+ * @param pool
+ * @param cbs
+ */
+void
+lua_thread_pool_prepare_callback (struct lua_thread_pool *pool, struct lua_callback_state *cbs);
+
+/**
+ * Restores state after lua_thread_pool_prepare_callback () usage
+ *
+ * @param cbs
+ */
+void
+lua_thread_pool_restore_callback (struct lua_callback_state *cbs);
 
 #endif /* LUA_THREAD_POOL_H_ */
 
-- 
2.39.5