feat: support kafka consumer for pubsub scenario

.github/workflows/build.yml

@@ -31,7 +31,7 @@ jobs:
           - linux_openresty_1_17
         test_dir:
           - t/plugin
-          - t/admin t/cli t/config-center-yaml t/control t/core t/debug t/discovery t/error_page t/misc
+          - t/admin t/cli t/config-center-yaml t/control t/core t/debug t/discovery t/error_page t/misc t/pubsub
           - t/node t/router t/script t/stream-node t/utils t/wasm t/xds-library t/xrpc
 
     runs-on: ${{ matrix.platform }}

.github/workflows/centos7-ci.yml

@@ -29,7 +29,7 @@ jobs:
       matrix:
         test_dir:
           - t/plugin
-          - t/admin t/cli t/config-center-yaml t/control t/core t/debug t/discovery t/error_page t/misc
+          - t/admin t/cli t/config-center-yaml t/control t/core t/debug t/discovery t/error_page t/misc t/pubsub
           - t/node t/router t/script t/stream-node t/utils t/wasm t/xds-library
 
     steps:
@@ -79,6 +79,8 @@ jobs:
 
     - name: Run other docker containers for test
       run: |
+        # generating SSL certificates for Kafka
+        keytool -genkeypair -keyalg RSA -dname "CN=127.0.0.1" -alias 127.0.0.1 -keystore ./ci/pod/kafka/kafka-server/selfsigned.jks -validity 365 -keysize 2048 -storepass changeit
         make ci-env-up
         ./ci/linux-ci-init-service.sh
 

Makefile

@@ -264,6 +264,9 @@ install: runtime
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix
 	$(ENV_INSTALL) apisix/*.lua $(ENV_INST_LUADIR)/apisix/
 
+	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/include/apisix/model
+	$(ENV_INSTALL) apisix/include/apisix/model/*.proto $(ENV_INST_LUADIR)/apisix/include/apisix/model/
+
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/admin
 	$(ENV_INSTALL) apisix/admin/*.lua $(ENV_INST_LUADIR)/apisix/admin/
 
@@ -340,6 +343,9 @@ install: runtime
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/zipkin
 	$(ENV_INSTALL) apisix/plugins/zipkin/*.lua $(ENV_INST_LUADIR)/apisix/plugins/zipkin/
 
+	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/pubsub
+	$(ENV_INSTALL) apisix/pubsub/*.lua $(ENV_INST_LUADIR)/apisix/pubsub/
+
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/ssl/router
 	$(ENV_INSTALL) apisix/ssl/router/*.lua $(ENV_INST_LUADIR)/apisix/ssl/router/
 

apisix/core.lua

@@ -52,4 +52,5 @@ return {
     tablepool   = require("tablepool"),
     resolver    = require("apisix.core.resolver"),
     os          = require("apisix.core.os"),
+    pubsub      = require("apisix.core.pubsub"),
 }

apisix/core/pubsub.lua

@@ -0,0 +1,167 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+
+--- Extensible framework to support publish-and-subscribe scenarios
+--
+-- @module core.pubsub
+
+local log          = require("apisix.core.log")
+local ws_server    = require("resty.websocket.server")
+local protoc       = require("protoc")
+local pb           = require("pb")
+local setmetatable = setmetatable
+local pcall        = pcall
+local pairs        = pairs
+
+protoc.reload()
+pb.option("int64_as_string")
+local pubsub_protoc = protoc.new()
+
+
+local _M = { version = 0.1 }
+local mt = { __index = _M }
+
+
+---
+-- Create pubsub module instance
+--
+-- @function core.pubsub.new
+-- @treturn pubsub module instance
+-- @treturn string|nil error message if present
+-- @usage
+-- local pubsub, err = core.pubsub.new()
+function _M.new()
+    -- compile the protobuf file on initial load module
+    -- ensure that each worker is loaded once
+    if not pubsub_protoc.loaded["pubsub.proto"] then
+        pubsub_protoc:addpath("apisix/include/apisix/model")
+        local ok, err = pcall(pubsub_protoc.loadfile, pubsub_protoc, "pubsub.proto")
+        if not ok then
+            pubsub_protoc:reset()
+            return nil, "failed to load pubsub protocol: "..err
+        end
+    end
+
+    local ws, err = ws_server:new()
+    if not ws then
+        return nil, err
+    end
+
+    local obj = setmetatable({
+        ws_server = ws,
+        cmd_handler = {},
+    }, mt)
+
+    return obj
+end
+
+
+---
+-- Add command callbacks to pubsub module instances
+--
+-- The callback function prototype: function (params)
+-- The params in the parameters contain the data defined in the requested command.
+-- Its first return value is the data, which needs to contain the data needed for
+-- the particular resp, returns nil if an error exists.
+-- Its second return value is a string type error message, no need to return when
+-- no error exists.
+--
+-- @function core.pubsub.on
+-- @usage
+-- pubsub:on(command, function (params)
+--     return data, err
+-- end)
+function _M.on(self, command, handler)
+    self.cmd_handler[command] = handler
+end
+
+
+---
+-- Put the pubsub instance into an event loop, waiting to process client commands
+--
+-- @function core.pubsub.wait
+-- @treturn string|nil error message if present, will terminate the event loop
+-- @usage
+-- local err = pubsub:wait()
+function _M.wait(self)
+    local ws = self.ws_server
+    while true do
+        -- read raw data frames from websocket connection
+        local raw_data, raw_type, err = ws:recv_frame()
+        if err then
+            -- terminate the event loop when a fatal error occurs
+            if ws.fatal then
+                ws:send_close()
+                return "websocket server: "..err
+            end
+
+            -- skip this loop for non-fatal errors
+            log.error("failed to receive websocket frame: "..err)
+            goto continue
+        end
+
+        -- handle client close connection
+        if raw_type == "close" then
+            ws:send_close()
+            return
+        end
+
+        -- the pub-sub messages use binary, if the message is not
+        -- binary, skip this message
+        if raw_type ~= "binary" then
+            goto continue
+        end
+
+        local data = pb.decode("PubSubReq", raw_data)
+        local sequence = data.sequence
+
+        -- call command handler to generate response data
+        for key, value in pairs(data) do
+            -- There are sequence and command properties in the data,
+            -- select the handler according to the command value.
+            if key ~= "sequence" then
+                local handler = self.cmd_handler[key]
+                if not handler then
+                    log.error("callback handler not registered for the",
+                        " this command, command: ", key)
+                    goto continue
+                end
+
+                local resp, err = handler(value)
+                if not resp then
+                    ws:send_binary(pb.encode("PubSubResp", {
+                        sequence = sequence,
+                        error_resp = {
+                            code = 0,
+                            message = err,
+                        },
+                    }))
+                    goto continue
+                end
+
+                -- write back the sequence
+                resp.sequence = sequence
+                ws:send_binary(pb.encode("PubSubResp", resp))
+            end
+        end
+
+        ::continue::
+    end
+end
+
+
+return _M

apisix/include/apisix/model/pubsub.proto

@@ -0,0 +1,117 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+syntax = "proto3";
+
+option java_package = "org.apache.apisix.api.pubsub";
+option java_outer_classname = "PubSubProto";
+option java_multiple_files = true;
+option go_package = "github.com/apache/apisix/api/pubsub;pubsub";
+
+/**
+ * Get the offset of the specified topic partition from Apache Kafka.
+ */
+message CmdKafkaListOffset {
+    string topic = 1;
+    int32 partition = 2;
+    int64 timestamp = 3;
+}
+
+/**
+ * Fetch messages of the specified topic partition from Apache Kafka.
+ */
+message CmdKafkaFetch {
+    string topic = 1;
+    int32 partition = 2;
+    int64 offset = 3;
+}
+
+/**
+ * Client request definition for pubsub scenarios
+ *
+ * The sequence field is used to associate requests and responses.
+ * Apache APISIX will set a consistent sequence for the associated
+ * requests and responses, and the client can explicitly know the
+ * response corresponding to any of the requests.
+ *
+ * The req field is the command data sent by the client, and its
+ * type will be chosen from any of the lists in the definition.
+ *
+ * Field numbers 1 to 30 in the definition are used to define basic
+ * information and future extensions, and numbers after 30 are used
+ * to define commands.
+ */
+message PubSubReq {
+    int64 sequence = 1;
+    oneof req {
+        CmdKafkaFetch      cmd_kafka_fetch       = 31;
+        CmdKafkaListOffset cmd_kafka_list_offset = 32;
+    };
+}
+
+/**
+ * The definition of a message in Kafka with the current message
+ * offset, production timestamp, Key, and message content.
+ */
+message KafkaMessage {
+    int64 offset = 1;
+    int64 timestamp = 2;
+    bytes key = 3;
+    bytes value = 4;
+}
+
+/**
+ * The response body of the service when an error occurs,
+ * containing the error code and the error message.
+ */
+message ErrorResp {
+    int32 code = 1;
+    string message = 2;
+}
+
+/**
+ * The response of Fetch messages from Apache Kafka.
+ */
+message KafkaFetchResp {
+    repeated KafkaMessage messages = 1;
+}
+
+/**
+ * The response of list offset from Apache Kafka.
+ */
+message KafkaListOffsetResp {
+    int64 offset = 1;
+}
+
+/**
+ * Server response definition for pubsub scenarios
+ *
+ * The sequence field will be the same as the value in the
+ * request, which is used to associate the associated request
+ * and response.
+ *
+ * The resp field is the response data sent by the server, and
+ * its type will be chosen from any of the lists in the definition.
+ */
+message PubSubResp {
+    int64 sequence = 1;
+    oneof resp {
+        ErrorResp error_resp = 31;
+        KafkaFetchResp kafka_fetch_resp = 32;
+        KafkaListOffsetResp kafka_list_offset_resp = 33;
+    };
+}

apisix/init.lua

@@ -42,6 +42,7 @@ local xrpc            = require("apisix.stream.xrpc")
 local ctxdump         = require("resty.ctxdump")
 local ngx_balancer    = require("ngx.balancer")
 local debug           = require("apisix.debug")
+local pubsub_kafka    = require("apisix.pubsub.kafka")
 local ngx             = ngx
 local get_method      = ngx.req.get_method
 local ngx_exit        = ngx.exit
@@ -504,6 +505,11 @@ function _M.http_access_phase()
         api_ctx.upstream_scheme = "grpc"
     end
 
+    -- load balancer is not required by kafka upstream
+    if api_ctx.matched_upstream and api_ctx.matched_upstream.scheme == "kafka" then
+        return pubsub_kafka.access(api_ctx)
+    end
+
     local code, err = set_upstream(route, api_ctx)
     if code then
         core.log.error("failed to set upstream: ", err)