feat: plugin grpc-transcode supports grpc deadline (#1149)

lua/apisix/plugins/grpc-transcode.lua

@@ -53,8 +53,19 @@ local schema = {
     type = "object",
     properties = {
         proto_id  = schema_def.id_schema,
-        service   = { type = "string" },
-        method    = { type = "string" },
+        service = {
+            description = "the grpc service name",
+            type        = "string"
+        },
+        method = {
+            description = "the method name in the grpc service.",
+            type    = "string"
+        },
+        deadline = {
+            description = "deadline for grpc, millisecond",
+            type        = "number",
+            default     = 0
+        },
         pb_option = { type = "array",
                       items = { type="string", anyOf = pb_option_def },
                       minItems = 1,
@@ -63,6 +74,17 @@ local schema = {
     additionalProperties = true }
 }
 
+local status_rel = {
+    ["3"] = 400,
+    ["4"] = 504,
+    ["5"] = 404,
+    ["7"] = 403,
+    ["11"] = 416,
+    ["12"] = 501,
+    ["13"] = 500,
+    ["14"] = 503,
+}
+
 local _M = {
     version = 0.1,
     priority = 506,
@@ -102,7 +124,7 @@ function _M.access(conf, ctx)
     end
 
     local ok, err = request(proto_obj, conf.service,
-                            conf.method, conf.pb_option)
+                            conf.method, conf.pb_option, conf.deadline)
     if not ok then
         core.log.error("transform request error: ", err)
         return
@@ -119,6 +141,18 @@ function _M.header_filter(conf, ctx)
 
     ngx.header["Content-Type"] = "application/json"
     ngx.header["Trailer"] = {"grpc-status", "grpc-message"}
+
+    local headers = ngx.resp.get_headers()
+    if headers["grpc-status"] ~= nil and headers["grpc-status"] ~= "0" then
+        local http_status = status_rel[headers["grpc-status"]]
+        if http_status ~= nil then
+            ngx.status = http_status
+        else
+            ngx.status = 599
+        end
+        return
+    end
+
 end
 
 

lua/apisix/plugins/grpc-transcode/request.lua

@@ -23,8 +23,9 @@ local ngx    = ngx
 local string = string
 local table  = table
 local ipairs = ipairs
+local tonumber = tonumber
 
-return function (proto, service, method, pb_option, default_values)
+return function (proto, service, method, pb_option, deadline, default_values)
     core.log.info("proto: ", core.json.delay_encode(proto, true))
     local m = util.find_method(proto, service, method)
     if not m then
@@ -62,5 +63,11 @@ return function (proto, service, method, pb_option, default_values)
     ngx.req.set_uri("/" .. service .. "/" .. method, false)
     ngx.req.set_uri_args({})
     ngx.req.set_body_data(message)
+
+    local dl = tonumber(deadline)
+    if dl~= nil and dl > 0 then
+        ngx.req.set_header("grpc-timeout",  dl .. "m")
+    end
+
     return true
 end

lua/apisix/plugins/grpc-transcode/response.lua

@@ -41,33 +41,33 @@ return function(proto, service, method, pb_option)
         ngx.arg[1] = nil
     end
 
-    if not eof then
-        return
-    end
 
-    ngx.ctx.buffered = nil
-    local buffer = table.concat(buffered)
-    if not ngx.req.get_headers()["X-Grpc-Web"] then
-        buffer = string.sub(buffer, 6)
-    end
+    if eof then
+        ngx.ctx.buffered = nil
+        local buffer = table.concat(buffered)
+        if not ngx.req.get_headers()["X-Grpc-Web"] then
+            buffer = string.sub(buffer, 6)
+        end
 
-    if pb_option then
-        for _, opt in ipairs(pb_option) do
-            pb.option(opt)
+        if pb_option then
+            for _, opt in ipairs(pb_option) do
+                pb.option(opt)
+            end
         end
-    end
 
-    local decoded = pb.decode(m.output_type, buffer)
-    if not decoded then
-        ngx.arg[1] = "failed to decode response data by protobuf"
-        return
-    end
+        local decoded = pb.decode(m.output_type, buffer)
+        if not decoded then
+            ngx.arg[1] = "failed to decode response data by protobuf"
+            return "failed to decode response data by protobuf"
+        end
+
+        local response, err = core.json.encode(decoded)
+        if not response then
+            core.log.error("failed to call json_encode data: ", err)
+            response = "failed to json_encode response body"
+        end
 
-    local response, err = core.json.encode(decoded)
-    if not response then
-        core.log.error("failed to call json_encode data: ", err)
-        response = "failed to json_encode response body"
+        ngx.arg[1] = response
     end
 
-    ngx.arg[1] = response
 end

t/plugin/grpc-transcode.t

@@ -232,7 +232,9 @@ Connection refused) while connecting to upstream
                       service Greeter {
                           rpc SayHello (HelloRequest) returns (HelloReply) {}
                           rpc Plus (PlusRequest) returns (PlusReply) {}
+                          rpc SayHelloAfterDelay (HelloRequest) returns (HelloReply) {}
                       }
+
                       message HelloRequest {
                           string name = 1;
                       }
@@ -324,3 +326,106 @@ GET /grpc_plus?a=1&b=2251799813685260
 qr/\{"result":"#2251799813685261"\}/
 --- no_error_log
 [error]
+
+
+
+=== TEST 11: set route3 deadline nodelay
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/3',
+                ngx.HTTP_PUT,
+                [[{
+                    "methods": ["GET"],
+                    "uri": "/grpc_deadline",
+                    "service_protocol": "grpc",
+                    "plugins": {
+                        "grpc-transcode": {
+                            "proto_id": "1",
+                            "service": "helloworld.Greeter",
+                            "method": "SayHello",
+                            "deadline": 500
+                        }
+                    },
+                    "upstream": {
+                        "type": "roundrobin",
+                        "nodes": {
+                            "127.0.0.1:50051": 1
+                        }
+                    }
+                }]]
+            )
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 12: hit route
+--- request
+GET /grpc_deadline?name=apisix
+--- response_body eval
+qr/\{"message":"Hello apisix"\}/
+--- no_error_log
+[error]
+
+
+
+=== TEST 13: set route4 deadline delay
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/4',
+                ngx.HTTP_PUT,
+                [[{
+                    "methods": ["GET"],
+                    "uri": "/grpc_delay",
+                    "service_protocol": "grpc",
+                    "plugins": {
+                        "grpc-transcode": {
+                            "proto_id": "1",
+                            "service": "helloworld.Greeter",
+                            "method": "SayHelloAfterDelay",
+                            "deadline": 500
+                        }
+                    },
+                    "upstream": {
+                        "type": "roundrobin",
+                        "nodes": {
+                            "127.0.0.1:50051": 1
+                        }
+                    }
+                }]]
+            )
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 14: hit route
+--- request
+GET /grpc_delay?name=apisix
+--- error_code: 504
+
+