[Feature] Add per-task lua cache to reuse 'heavy' objects

diff --git a/src/libserver/task.c b/src/libserver/task.c
index dc3198609892a8c87a2f512e97549612127a46a3..886e1b0040399bbcc8a93577811dd33676ea8570 100644 (file)
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -129,6 +129,7 @@ rspamd_task_new (struct rspamd_worker *worker, struct rspamd_config *cfg)
 
        new_task->message_id = new_task->queue_id = "undef";
        new_task->messages = ucl_object_typed_new (UCL_OBJECT);
+       new_task->lua_cache = g_hash_table_new (rspamd_str_hash, rspamd_str_equal);
 
        return new_task;
 }
@@ -192,6 +193,9 @@ rspamd_task_free (struct rspamd_task *task)
        struct rspamd_mime_part *p;
        struct rspamd_mime_text_part *tp;
        struct rspamd_email_address *addr;
+       GHashTableIter it;
+       gpointer k, v;
+       gint lua_ref;
        guint i;
 
        if (task) {
@@ -236,6 +240,7 @@ rspamd_task_free (struct rspamd_task *task)
                }
 
                ucl_object_unref (task->messages);
+               rspamd_re_cache_runtime_destroy (task->re_rt);
 
                if (task->http_conn != NULL) {
                        rspamd_http_connection_reset (task->http_conn);
@@ -271,7 +276,18 @@ rspamd_task_free (struct rspamd_task *task)
                }
 
                if (task->cfg) {
-                       rspamd_re_cache_runtime_destroy (task->re_rt);
+                       if (task->lua_cache) {
+                               g_hash_table_iter_init (&it, task->lua_cache);
+
+                               while (g_hash_table_iter_next (&it, &k, &v)) {
+                                       lua_ref = GPOINTER_TO_INT (v);
+                                       luaL_unref (task->cfg->lua_state,
+                                                       LUA_REGISTRYINDEX, lua_ref);
+                               }
+
+                               g_hash_table_unref (task->lua_cache);
+                       }
+
                        REF_RELEASE (task->cfg);
                }
 
@@ -1489,7 +1505,7 @@ rspamd_task_profile_set (struct rspamd_task *task, const gchar *key,
        tbl = rspamd_mempool_get_variable (task->task_pool, "profile");
 
        if (tbl == NULL) {
-               tbl = g_hash_table_new (g_str_hash, g_str_equal);
+               tbl = g_hash_table_new (rspamd_str_hash, rspamd_str_equal);
                rspamd_mempool_set_variable (task->task_pool, "profile", tbl,
                                (rspamd_mempool_destruct_t)g_hash_table_unref);
        }

diff --git a/src/libserver/task.h b/src/libserver/task.h
index 1779310bb7340f2b867cdf3af4d0a2b612095012..c6dff0ebc9f559a627a3bbd9a3daca2cff9ea500 100644 (file)
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -154,6 +154,7 @@ struct rspamd_task {
        GHashTable *raw_headers;                                                /**< list of raw headers                                                        */
        GHashTable *results;                                                    /**< hash table of metric_result indexed by
                                                                                                         *    metric's name                                                                     */
+       GHashTable *lua_cache;                                                  /**< cache of lua objects                                                       */
        GPtrArray *tokens;                                                              /**< statistics tokens */
 
        GPtrArray *rcpt_mime;

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index d7fabd213a8c7ee9e3dbe7aa6ddaa67399d40ee5..eb9ddbc024e15e154913da802afeabf490f46818 100644 (file)
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -867,6 +867,43 @@ lua_check_text (lua_State * L, gint pos)
        return ud ? (struct rspamd_lua_text *)ud : NULL;
 }
 
+static void
+lua_task_set_cached (lua_State *L, struct rspamd_task *task, const gchar *key,
+               gint pos)
+{
+       gpointer elt;
+       gint lua_ref;
+
+       lua_pushvalue (L, pos);
+
+       elt = g_hash_table_lookup (task->lua_cache, key);
+
+       if (G_UNLIKELY (elt != NULL)) {
+               /* Unref previous value */
+               lua_ref = GPOINTER_TO_INT (elt);
+               luaL_unref (L, LUA_REGISTRYINDEX, lua_ref);
+       }
+
+       lua_ref = luaL_ref (L, LUA_REGISTRYINDEX);
+       g_hash_table_insert (task->lua_cache, (void *)key, GINT_TO_POINTER (lua_ref));
+}
+
+static gboolean
+lua_task_get_cached (lua_State *L, struct rspamd_task *task, const gchar *key)
+{
+       gpointer elt;
+
+       elt = g_hash_table_lookup (task->lua_cache, key);
+
+       if (elt != NULL) {
+               lua_rawgeti (L, LUA_REGISTRYINDEX, GPOINTER_TO_INT (elt));
+
+               return TRUE;
+       }
+
+       return FALSE;
+}
+
 /* Task methods */
 static int
 lua_task_process_message (lua_State *L)
@@ -1157,19 +1194,32 @@ lua_task_get_urls (lua_State * L)
                        need_emails = lua_toboolean (L, 2);
                }
 
-               sz = g_hash_table_size (task->urls);
-
                if (need_emails) {
-                       sz += g_hash_table_size (task->emails);
+                       if (!lua_task_get_cached (L, task, "emails+urls")) {
+                               sz = g_hash_table_size (task->urls);
+                               sz += g_hash_table_size (task->emails);
+
+                               lua_createtable (L, sz, 0);
+                               cb.i = 1;
+                               cb.L = L;
+                               g_hash_table_foreach (task->urls, lua_tree_url_callback, &cb);
+                               g_hash_table_foreach (task->emails, lua_tree_url_callback, &cb);
+
+                               lua_task_set_cached (L, task, "emails+urls", -1);
+                       }
+
                }
+               else {
+                       if (!lua_task_get_cached (L, task, "urls")) {
+                               sz = g_hash_table_size (task->urls);
 
-               lua_createtable (L, sz, 0);
-               cb.i = 1;
-               cb.L = L;
-               g_hash_table_foreach (task->urls, lua_tree_url_callback, &cb);
+                               lua_createtable (L, sz, 0);
+                               cb.i = 1;
+                               cb.L = L;
+                               g_hash_table_foreach (task->urls, lua_tree_url_callback, &cb);
 
-               if (need_emails) {
-                       g_hash_table_foreach (task->emails, lua_tree_url_callback, &cb);
+                               lua_task_set_cached (L, task, "urls", -1);
+                       }
                }
        }
        else {
@@ -1283,15 +1333,20 @@ lua_task_get_text_parts (lua_State * L)
        struct rspamd_mime_text_part *part, **ppart;
 
        if (task != NULL) {
-               lua_createtable (L, task->text_parts->len, 0);
-
-               for (i = 0; i < task->text_parts->len; i ++) {
-                       part = g_ptr_array_index (task->text_parts, i);
-                       ppart = lua_newuserdata (L, sizeof (struct rspamd_mime_text_part *));
-                       *ppart = part;
-                       rspamd_lua_setclass (L, "rspamd{textpart}", -1);
-                       /* Make it array */
-                       lua_rawseti (L, -2, i + 1);
+
+               if (!lua_task_get_cached (L, task, "text_parts")) {
+                       lua_createtable (L, task->text_parts->len, 0);
+
+                       for (i = 0; i < task->text_parts->len; i ++) {
+                               part = g_ptr_array_index (task->text_parts, i);
+                               ppart = lua_newuserdata (L, sizeof (struct rspamd_mime_text_part *));
+                               *ppart = part;
+                               rspamd_lua_setclass (L, "rspamd{textpart}", -1);
+                               /* Make it array */
+                               lua_rawseti (L, -2, i + 1);
+                       }
+
+                       lua_task_set_cached (L, task, "text_parts", -1);
                }
        }
        else {
@@ -1309,15 +1364,19 @@ lua_task_get_parts (lua_State * L)
        struct rspamd_mime_part *part, **ppart;
 
        if (task != NULL) {
-               lua_createtable (L, task->parts->len, 0);
-
-               for (i = 0; i < task->parts->len; i ++) {
-                       part = g_ptr_array_index (task->parts, i);
-                       ppart = lua_newuserdata (L, sizeof (struct rspamd_mime_part *));
-                       *ppart = part;
-                       rspamd_lua_setclass (L, "rspamd{mimepart}", -1);
-                       /* Make it array */
-                       lua_rawseti (L, -2, i + 1);
+               if (!lua_task_get_cached (L, task, "mime_parts")) {
+                       lua_createtable (L, task->parts->len, 0);
+
+                       for (i = 0; i < task->parts->len; i ++) {
+                               part = g_ptr_array_index (task->parts, i);
+                               ppart = lua_newuserdata (L, sizeof (struct rspamd_mime_part *));
+                               *ppart = part;
+                               rspamd_lua_setclass (L, "rspamd{mimepart}", -1);
+                               /* Make it array */
+                               lua_rawseti (L, -2, i + 1);
+                       }
+
+                       lua_task_set_cached (L, task, "mime_parts", -1);
                }
        }
        else {
@@ -1543,70 +1602,77 @@ lua_task_get_received_headers (lua_State * L)
        guint i, k = 1;
 
        if (task) {
-               lua_createtable (L, task->received->len, 0);
+               if (!lua_task_get_cached (L, task, "received")) {
+                       lua_createtable (L, task->received->len, 0);
 
-               for (i = 0; i < task->received->len; i ++) {
-                       rh = g_ptr_array_index (task->received, i);
+                       for (i = 0; i < task->received->len; i ++) {
+                               rh = g_ptr_array_index (task->received, i);
 
-                       lua_createtable (L, 0, 9);
-                       rspamd_lua_table_set (L, "raw", rh->hdr->decoded);
+                               lua_createtable (L, 0, 9);
 
-                       if (G_UNLIKELY (rh->from_ip == NULL &&
-                                       rh->real_ip == NULL &&
-                                       rh->real_hostname == NULL &&
-                                       rh->by_hostname == NULL && rh->timestamp == 0 &&
-                                       rh->for_mbox == NULL)) {
-                               lua_rawseti (L, -2, k ++);
+                               if (rh->hdr && rh->hdr->decoded) {
+                                       rspamd_lua_table_set (L, "raw", rh->hdr->decoded);
+                               }
 
-                               continue;
-                       }
+                               if (G_UNLIKELY (rh->from_ip == NULL &&
+                                               rh->real_ip == NULL &&
+                                               rh->real_hostname == NULL &&
+                                               rh->by_hostname == NULL && rh->timestamp == 0 &&
+                                               rh->for_mbox == NULL)) {
+                                       lua_rawseti (L, -2, k ++);
 
-                       rspamd_lua_table_set (L, "from_hostname", rh->from_hostname);
-                       rspamd_lua_table_set (L, "from_ip", rh->from_ip);
-                       rspamd_lua_table_set (L, "real_hostname", rh->real_hostname);
-                       lua_pushstring (L, "real_ip");
-                       rspamd_lua_ip_push (L, rh->addr);
-                       lua_settable (L, -3);
-                       lua_pushstring (L, "proto");
+                                       continue;
+                               }
 
-                       switch (rh->type) {
-                       case RSPAMD_RECEIVED_SMTP:
-                               proto = "smtp";
-                               break;
-                       case RSPAMD_RECEIVED_ESMTP:
-                               proto = "esmtp";
-                               break;
-                       case RSPAMD_RECEIVED_ESMTPS:
-                               proto = "esmtps";
-                               break;
-                       case RSPAMD_RECEIVED_ESMTPA:
-                               proto = "esmtpa";
-                               break;
-                       case RSPAMD_RECEIVED_ESMTPSA:
-                               proto = "esmtpsa";
-                               break;
-                       case RSPAMD_RECEIVED_LMTP:
-                               proto = "lmtp";
-                               break;
-                       case RSPAMD_RECEIVED_IMAP:
-                               proto = "imap";
-                               break;
-                       case RSPAMD_RECEIVED_UNKNOWN:
-                       default:
-                               proto = "unknown";
-                               break;
-                       }
+                               rspamd_lua_table_set (L, "from_hostname", rh->from_hostname);
+                               rspamd_lua_table_set (L, "from_ip", rh->from_ip);
+                               rspamd_lua_table_set (L, "real_hostname", rh->real_hostname);
+                               lua_pushstring (L, "real_ip");
+                               rspamd_lua_ip_push (L, rh->addr);
+                               lua_settable (L, -3);
+                               lua_pushstring (L, "proto");
 
-                       lua_pushstring (L, proto);
-                       lua_settable (L, -3);
+                               switch (rh->type) {
+                               case RSPAMD_RECEIVED_SMTP:
+                                       proto = "smtp";
+                                       break;
+                               case RSPAMD_RECEIVED_ESMTP:
+                                       proto = "esmtp";
+                                       break;
+                               case RSPAMD_RECEIVED_ESMTPS:
+                                       proto = "esmtps";
+                                       break;
+                               case RSPAMD_RECEIVED_ESMTPA:
+                                       proto = "esmtpa";
+                                       break;
+                               case RSPAMD_RECEIVED_ESMTPSA:
+                                       proto = "esmtpsa";
+                                       break;
+                               case RSPAMD_RECEIVED_LMTP:
+                                       proto = "lmtp";
+                                       break;
+                               case RSPAMD_RECEIVED_IMAP:
+                                       proto = "imap";
+                                       break;
+                               case RSPAMD_RECEIVED_UNKNOWN:
+                               default:
+                                       proto = "unknown";
+                                       break;
+                               }
 
-                       lua_pushstring (L, "timestamp");
-                       lua_pushnumber (L, rh->timestamp);
-                       lua_settable (L, -3);
+                               lua_pushstring (L, proto);
+                               lua_settable (L, -3);
 
-                       rspamd_lua_table_set (L, "by_hostname", rh->by_hostname);
-                       rspamd_lua_table_set (L, "for", rh->for_mbox);
-                       lua_rawseti (L, -2, k ++);
+                               lua_pushstring (L, "timestamp");
+                               lua_pushnumber (L, rh->timestamp);
+                               lua_settable (L, -3);
+
+                               rspamd_lua_table_set (L, "by_hostname", rh->by_hostname);
+                               rspamd_lua_table_set (L, "for", rh->for_mbox);
+                               lua_rawseti (L, -2, k ++);
+                       }
+
+                       lua_task_set_cached (L, task, "received", -1);
                }
        }
        else {
@@ -2248,17 +2314,21 @@ lua_task_get_images (lua_State *L)
        struct rspamd_image **pimg;
 
        if (task) {
-               lua_newtable (L);
+               if (!lua_task_get_cached (L, task, "images")) {
+                       lua_createtable (L, task->parts->len, 0);
 
-               for (i = 0; i < task->parts->len; i ++) {
-                       part = g_ptr_array_index (task->parts, i);
+                       for (i = 0; i < task->parts->len; i ++) {
+                               part = g_ptr_array_index (task->parts, i);
 
-                       if (part->flags & RSPAMD_MIME_PART_IMAGE) {
-                               pimg = lua_newuserdata (L, sizeof (struct rspamd_image *));
-                               rspamd_lua_setclass (L, "rspamd{image}", -1);
-                               *pimg = part->specific.img;
-                               lua_rawseti (L, -2, ++nelt);
+                               if (part->flags & RSPAMD_MIME_PART_IMAGE) {
+                                       pimg = lua_newuserdata (L, sizeof (struct rspamd_image *));
+                                       rspamd_lua_setclass (L, "rspamd{image}", -1);
+                                       *pimg = part->specific.img;
+                                       lua_rawseti (L, -2, ++nelt);
+                               }
                        }
+
+                       lua_task_set_cached (L, task, "images", -1);
                }
        }
        else {
@@ -2277,17 +2347,21 @@ lua_task_get_archives (lua_State *L)
        struct rspamd_archive **parch;
 
        if (task) {
-               lua_createtable (L, task->parts->len, 0);
+               if (!lua_task_get_cached (L, task, "archives")) {
+                       lua_createtable (L, task->parts->len, 0);
 
-               for (i = 0; i < task->parts->len; i ++) {
-                       part = g_ptr_array_index (task->parts, i);
+                       for (i = 0; i < task->parts->len; i ++) {
+                               part = g_ptr_array_index (task->parts, i);
 
-                       if (part->flags & RSPAMD_MIME_PART_ARCHIVE) {
-                               parch = lua_newuserdata (L, sizeof (struct rspamd_archive *));
-                               rspamd_lua_setclass (L, "rspamd{archive}", -1);
-                               *parch = part->specific.arch;
-                               lua_rawseti (L, -2, ++nelt);
+                               if (part->flags & RSPAMD_MIME_PART_ARCHIVE) {
+                                       parch = lua_newuserdata (L, sizeof (struct rspamd_archive *));
+                                       rspamd_lua_setclass (L, "rspamd{archive}", -1);
+                                       *parch = part->specific.arch;
+                                       lua_rawseti (L, -2, ++nelt);
+                               }
                        }
+
+                       lua_task_set_cached (L, task, "archives", -1);
                }
        }
        else {