[Project] Fix Lua DNS code

2 files changed, 21 insertions, 14 deletions

diff --git a/src/lua/lua_dns.c b/src/lua/lua_dns.c
index fe4b7ff2e..da026a8a3 100644
--- a/src/lua/lua_dns.c
+++ b/src/lua/lua_dns.c
@@ -30,7 +30,7 @@ struct lua_rspamd_dns_cbdata {
 	struct thread_entry *thread;
 	struct rspamd_task *task;
 	struct rspamd_dns_resolver *resolver;
-	struct rspamd_async_watcher *w;
+	struct rspamd_symcache_item *item;
 	struct rspamd_async_session *s;
 };
 
@@ -50,8 +50,13 @@ lua_dns_request (lua_State *L)
 
 	/* Check arguments */
 	if (!rspamd_lua_parse_table_arguments (L, 1, &err,
-										   "*name=S;task=U{task};*type=S;forced=B;session=U{session};config=U{config}",
-										   &to_resolve, &task, &type_str, &forced, &session, &cfg)) {
+			"*name=S;task=U{task};*type=S;forced=B;session=U{session};config=U{config}",
+			&to_resolve,
+			&task,
+			&type_str,
+			&forced,
+			&session,
+			&cfg)) {
 
 		if (err) {
 			ret = luaL_error (L, "invalid arguments: %s", err->message);
@@ -132,8 +137,12 @@ lua_dns_request (lua_State *L)
 	if (ret) {
 		cbdata->thread = lua_thread_pool_get_running_entry (cfg->lua_thread_pool);
 		cbdata->s = session;
-		cbdata->w = rspamd_session_get_watcher (session);
-		rspamd_session_watcher_push (session);
+
+		if (task) {
+			cbdata->item = rspamd_symbols_cache_get_cur_item (task);
+			rspamd_symcache_item_async_inc (task, cbdata->item);
+		}
+
 		return lua_thread_yield (cbdata->thread, 0);
 	}
 	else {
@@ -166,8 +175,8 @@ lua_dns_callback (struct rdns_reply *reply, void *arg)
 
 	lua_thread_resume (cbdata->thread, 2);
 
-	if (cbdata->s) {
-		rspamd_session_watcher_pop (cbdata->s, cbdata->w);
+	if (cbdata->item) {
+		rspamd_symcache_item_async_dec_check (cbdata->task, cbdata->item);
 	}
 }
 
diff --git a/src/lua/lua_dns_resolver.c b/src/lua/lua_dns_resolver.c
index e7b6eeaab..dfdac1630 100644
--- a/src/lua/lua_dns_resolver.c
+++ b/src/lua/lua_dns_resolver.c
@@ -86,7 +86,7 @@ struct lua_dns_cbdata {
 	gint cbref;
 	gchar *to_resolve;
 	gchar *user_str;
-	struct rspamd_async_watcher *w;
+	struct rspamd_symcache_item *item;
 	struct rspamd_async_session *s;
 };
 
@@ -199,8 +199,8 @@ lua_dns_resolver_callback (struct rdns_reply *reply, gpointer arg)
 	luaL_unref (L, LUA_REGISTRYINDEX, cd->cbref);
 	lua_thread_pool_restore_callback (&cbs);
 
-	if (cd->s) {
-		rspamd_session_watcher_pop (cd->s, cd->w);
+	if (cd->item) {
+		rspamd_symcache_item_async_dec_check (cd->task, cd->item);
 	}
 
 	if (!cd->pool) {
@@ -432,8 +432,6 @@ lua_dns_resolver_resolve_common (lua_State *L,
 
 				if (session) {
 					cbdata->s = session;
-					cbdata->w = rspamd_session_get_watcher (session);
-					rspamd_session_watcher_push (session);
 				}
 			}
 			else {
@@ -457,8 +455,8 @@ lua_dns_resolver_resolve_common (lua_State *L,
 
 			if (ret) {
 				cbdata->s = session;
-				cbdata->w = rspamd_session_get_watcher (session);
-				rspamd_session_watcher_push (session);
+				cbdata->item = rspamd_symbols_cache_get_cur_item (task);
+				rspamd_symcache_item_async_inc (task, cbdata->item);
 				/* callback was set up */
 				lua_pushboolean (L, TRUE);
 			} else {