msb-core/openresty-ext/src/assembly/resources/openresty/nginx/luaext/customrouter.lua

   1 --[[
   2
   3     Copyright 2016 ZTE Corporation.
   4
   5     Licensed under the Apache License, Version 2.0 (the "License");
   6     you may not use this file except in compliance with the License.
   7     You may obtain a copy of the License at
   8
   9         http://www.apache.org/licenses/LICENSE-2.0
  10
  11     Unless required by applicable law or agreed to in writing, software
  12     distributed under the License is distributed on an "AS IS" BASIS,
  13     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14     See the License for the specific language governing permissions and
  15     limitations under the License.
  16
  17         Author: Zhaoxing Meng
  18         email: meng.zhaoxing1@zte.com.cn
  19
  20 ]]
  21 -- put red into the connection pool of size 100,
  22 -- with 10 seconds max idle time
  23 local function close_redis(red)
  24         if not red then
  25                 return
  26         end
  27         --release connection to the pool
  28         local pool_max_idle_time = 10000
  29         local pool_size = 100
  30         local ok, err = red:set_keepalive(pool_max_idle_time, pool_size)
  31         if not ok then
  32                 ngx.log(ngx.ERR, "set keepalive error:", err)
  33         end
  34         return
  35 end
  36
  37 local function query_ipurl_updatecache(red,key)
  38         local keyprefix = "msb:routing:custom:"..key
  39
  40
  41         local infokey=keyprefix..":info"
  42         -- first of all check whether the status is 1(enabled)
  43         local status = red:hget(infokey,"status")
  44         if not (status=="1")  then
  45                 ngx.log(ngx.WARN, key.."is disabled.status=", status)
  46                 return nil
  47         end
  48
  49         -- Try to get url for key
  50         local url, err = red:hget(infokey,"url")
  51         ngx.log(ngx.WARN, "==url:", url)
  52         if not url or url == ngx.null then
  53                 return nil
  54         end
  55
  56         -- Try to get ip:port for key
  57         local serverkey=keyprefix..":lb:server1"
  58         local server, err = red:hget(serverkey,"ip")..":"..red:hget(serverkey,"port")
  59         ngx.log(ngx.WARN, "==server:", server)
  60         if not server or server == ngx.null then
  61                 return nil
  62         end
  63
  64         -- get the local cache
  65         local cache = ngx.shared.ceryx
  66         local uri = ngx.var.uri
  67         -- Save found key to local cache for 5 seconds
  68         cache:set("custom:key:"..uri,key,5)
  69         cache:set("custom:server:"..uri,server,5)
  70         cache:set("custom:url:"..uri,url,5)
  71
  72         ngx.var.key = key
  73         ngx.var.server = server
  74         ngx.var.url = url
  75         return true
  76 end
  77
  78 local function query_allkeys_updatecache(red)
  79         -- Try to get all keys start with "msb:routing:custom:"
  80         local allkeys, err = red:keys("msb:routing:custom:*")
  81         if not allkeys or allkeys == ngx.null then
  82                 ngx.log(ngx.ERR,err)
  83                 return ""
  84         end
  85         local key_set={}
  86         for key, value in ipairs(allkeys) do
  87                 name = string.gsub(string.gsub(string.gsub(value,"msb:routing:custom:",""),":info",""),":lb:server1","")
  88                 key_set[name]=true
  89         end
  90         local key_table = {}
  91         local index = 1
  92         for key,_ in pairs(key_set) do
  93                 key_table[index] = key
  94                 index = index + 1
  95         end
  96         table.sort(key_table, function (a, b)
  97                         return a > b
  98                 end)
  99
 100         local servicenames = ""
 101         local delimiter = "<>"
 102         for i=1,#key_table do
 103                 servicenames=servicenames..key_table[i]..delimiter
 104         end
 105
 106         -- get the local cache
 107         local cache = ngx.shared.ceryx
 108         -- Save all keys to local cache for 30 seconds(0.5 minutes)
 109         cache:set("customrouter:allkeys", servicenames, 30)
 110         return servicenames;
 111 end
 112
 113 local function query_router_info()
 114         local uri = ngx.var.uri
 115         ngx.log(ngx.WARN, "==uri:", uri)
 116
 117         -- Check if key exists in local cache
 118         local cache = ngx.shared.ceryx
 119         local key, flags = cache:get("custom:key:"..uri)
 120         local server, flags = cache:get("custom:server:"..uri)
 121         local url, flags = cache:get("custom:url:"..uri)
 122         if key and server and url then
 123                 ngx.var.key = key
 124                 ngx.var.server = server
 125                 ngx.var.url = url
 126                 ngx.log(ngx.WARN, "==using custom cache:", key.."&&"..server.."&&"..url)
 127                 return
 128         end
 129
 130         local redis = require "resty.redis"
 131         local red = redis:new()
 132         red:set_timeout(1000) -- 1000 ms
 133         local redis_host = "127.0.0.1"
 134         local redis_port = 6379
 135         local res, err = red:connect(redis_host, redis_port)
 136
 137         -- Return if could not connect to Redis
 138         if not res then
 139                 ngx.log(ngx.ERR, "connect to redis error:", err)
 140                 return
 141         end
 142
 143         -- Check if all servicenames exists in local cache
 144         local servicenames, flags = cache:get("customrouter:allkeys")
 145         if servicenames then
 146                 ngx.log(ngx.WARN,"==get all keys from cache:",servicenames)
 147         else
 148                 servicenames = query_allkeys_updatecache(red)
 149         end
 150
 151         local delimiter = "<>"
 152         for key in string.gmatch(servicenames,"(.-)"..delimiter) do
 153                 ngx.log(ngx.WARN, "==key_table key:", key)
 154                 local from, to, err = ngx.re.find(uri, "^"..key.."(/(.*))?$", "jo")
 155                 if from then
 156                         ngx.log(ngx.WARN,"Matched! start-end:",from,"-",to)
 157                         local result = query_ipurl_updatecache(red,key)
 158                         if result then
 159                                 break
 160                         end
 161                 else
 162                         ngx.log(ngx.WARN,"not Matched")
 163                         if err then
 164                                 ngx.log(ngx.WARN,"ngx.re.find throw error: ",err)
 165                                 return
 166                         end
 167                 end
 168         end
 169
 170         return close_redis(red)
 171 end
 172
 173 local function rewrite_router_url()
 174         local server=ngx.var.server
 175         if server=="fallback" then
 176                 ngx.status = ngx.HTTP_NOT_FOUND
 177                 ngx.exit(ngx.status)
 178         end
 179         local url=ngx.var.url
 180         local key=ngx.var.key
 181         local rewriteduri = ngx.re.sub(ngx.var.uri, "^"..key.."(.*)", url.."$1", "o")
 182         ngx.log(ngx.WARN, "==rewrited uri:", rewriteduri)
 183         ngx.req.set_uri(rewriteduri)
 184 end
 185
 186 query_router_info()
 187 rewrite_router_url()