From a15851eb98b0629d65956897d1853705055c1028 Mon Sep 17 00:00:00 2001
From: alonbg <alon.blayergat@gmail.com>
Date: Wed, 7 Sep 2016 16:39:19 +0300
Subject: [PATCH] cross modules stream/http shared dict

---
 src/ngx_stream_lua_shdict.c | 73 +++++++++++++++++++++++++++++++++----
 1 file changed, 66 insertions(+), 7 deletions(-)

diff --git a/src/ngx_stream_lua_shdict.c b/src/ngx_stream_lua_shdict.c
index 99472ce4..fbb55351 100644
--- a/src/ngx_stream_lua_shdict.c
+++ b/src/ngx_stream_lua_shdict.c
@@ -13,6 +13,7 @@
 #include "ngx_stream_lua_shdict.h"
 #include "ngx_stream_lua_util.h"
 
+extern ngx_module_t ngx_http_lua_module;
 
 static int ngx_stream_lua_shdict_set(lua_State *L);
 static int ngx_stream_lua_shdict_safe_set(lua_State *L);
@@ -308,12 +309,66 @@ ngx_stream_lua_shdict_expire(ngx_stream_lua_shdict_ctx_t *ctx, ngx_uint_t n)
 void
 ngx_stream_lua_inject_shdict_api(ngx_stream_lua_main_conf_t *lmcf, lua_State *L)
 {
-    ngx_stream_lua_shdict_ctx_t  *ctx;
-    ngx_uint_t                    i;
-    ngx_shm_zone_t              **zone;
+    ngx_uint_t                       i;
+    ngx_array_t                      all_zones;
+    ngx_stream_lua_shdict_ctx_t     *ctx;
+    ngx_shm_zone_t                  **zone, *shm_zone;
+    ngx_pool_t                      *temp_pool;
+    ngx_list_part_t                 *part;
+    ngx_log_t                       *log;
+
+    log = &lmcf->cycle->log;
+
+    /* place http and stream zones in a single array */
+    temp_pool = ngx_create_pool(NGX_DEFAULT_POOL_SIZE, log);
+
+    if (temp_pool == NULL) {
+        ngx_log_error(NGX_LOG_ERR, log, 1,
+                      "error: pool for zone array not allocated");
+        return;
+    }
+
+    if (ngx_array_init(&all_zones, temp_pool, 2,
+                       sizeof(ngx_shm_zone_t *)) != NGX_OK)
+    {
+        ngx_destroy_pool(temp_pool);
+        ngx_log_error(NGX_LOG_ERR, log, 1,
+                      "error: zone array not allocated");
+        return;
+    }
+
+    part = &lmcf->cycle->shared_memory.part;
+    shm_zone = part->elts;
+
+    for (i = 0; /* void */ ; i++) {
+
+        if (i >= part->nelts) {
+            if (part->next == NULL) {
+                break;
+            }
+            part = part->next;
+            shm_zone = part->elts;
+            i = 0;
+        }
+
+        if (&ngx_http_lua_module == shm_zone->tag ||
+            &ngx_stream_lua_module == shm_zone->tag )
+        {
+            zone = ngx_array_push(&all_zones);
+
+            if (zone == NULL) {
+                ngx_destroy_pool(temp_pool);
+                ngx_log_error(NGX_LOG_ERR, log, 1,
+                              "error: zone pointer not allocated");
+                return;
+            }
 
-    if (lmcf->shm_zones != NULL) {
-        lua_createtable(L, 0, lmcf->shm_zones->nelts /* nrec */);
+            *zone = shm_zone;
+        }
+    }
+
+    if (all_zones.nelts > 0) {
+        lua_createtable(L, 0, all_zones.nelts /* nrec */);
                 /* ngx.shared */
 
         lua_createtable(L, 0 /* narr */, 13 /* nrec */); /* shared mt */
@@ -357,11 +412,14 @@ ngx_stream_lua_inject_shdict_api(ngx_stream_lua_main_conf_t *lmcf, lua_State *L)
         lua_pushvalue(L, -1); /* shared mt mt */
         lua_setfield(L, -2, "__index"); /* shared mt */
 
-        zone = lmcf->shm_zones->elts;
+        zone = all_zones.elts;
 
-        for (i = 0; i < lmcf->shm_zones->nelts; i++) {
+        for (i = 0; i < all_zones.nelts; i++) {
             ctx = zone[i]->data;
 
+            dd("injecting shared dict %.*s",
+               (int) ctx->name.len, ctx->name.data);
+
             lua_pushlstring(L, (char *) ctx->name.data, ctx->name.len);
                 /* shared mt key */
 
@@ -381,6 +439,7 @@ ngx_stream_lua_inject_shdict_api(ngx_stream_lua_main_conf_t *lmcf, lua_State *L)
     }
 
     lua_setfield(L, -2, "shared");
+    ngx_destroy_pool(temp_pool);
 }