update rpc_module with rpc_filter

CMakeLists_Headers.txt

@@ -15,6 +15,7 @@ set(SRC_HEADERS
 	src/module/rpc_span_policies.h
 	src/module/rpc_module_span.h
 	src/module/rpc_module.h
+	src/module/rpc_filter.h
 	src/rpc_basic.h
 	src/rpc_buffer.h
 	src/rpc_client.h

src/message/rpc_message.h

@@ -21,7 +21,7 @@
 #include <string>
 #include <workflow/ProtocolMessage.h>
 #include "rpc_basic.h"
-#include "rpc_module.h"
+#include "rpc_filter.h"
 
 namespace srpc
 {

src/module/rpc_filter.h

@@ -0,0 +1,99 @@
+/*
+  Copyright (c) 2021 Sogou, Inc.
+
+  Licensed 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.
+*/
+
+#ifndef __RPC_FILTER_H__
+#define __RPC_FILTER_H__
+
+#include <map>
+#include <vector>
+#include <string>
+#include "workflow/WFTask.h"
+#include "workflow/WFTaskFactory.h"
+
+namespace srpc
+{
+
+using RPCModuleData = std::map<std::string, std::string>;
+using RPCLogVector = std::vector<std::pair<std::string, std::string>>;
+
+static RPCModuleData global_empty_map;
+
+class RPCFilter
+{
+public:
+	SubTask *create_filter_task(const RPCModuleData& data)
+	{
+		if (this->filter(const_cast<RPCModuleData&>(data)))
+			return this->create(const_cast<RPCModuleData&>(data));
+
+		return WFTaskFactory::create_empty_task();
+	}
+
+private:
+	virtual SubTask *create(RPCModuleData& data) = 0;
+
+	virtual bool filter(RPCModuleData& data) = 0;
+
+public:
+	RPCFilter(int module_type, const std::string name) :
+		name(std::move(name))
+	{
+		this->module_type = module_type;
+	}
+
+	const std::string& get_name() const { return this->name; }
+	int get_module_type() const { return this->module_type; }
+
+	virtual bool client_begin(SubTask *task, const RPCModuleData& data)
+	{
+		return true;
+	}
+
+	virtual bool server_begin(SubTask *task, const RPCModuleData& data)
+	{
+		return true;
+	}
+
+	virtual bool client_end(SubTask *task, const RPCModuleData& data)
+	{
+		return true;
+	}
+
+	virtual bool server_end(SubTask *task, const RPCModuleData& data)
+	{
+		return true;
+	}
+
+private:
+	std::string name;
+	int module_type;
+};
+
+/*
+class RPCMonitorPrometheus : public RPCFilter
+{
+public:
+	bool filter_finish()
+	{
+		SRPCGlobal->get_instance()->get_prometheus()->write();
+	}
+};
+*/
+
+} // end namespace srpc
+
+#endif
+

src/module/rpc_module.h

@@ -17,49 +17,111 @@
 #ifndef __RPC_MODULE_H__
 #define __RPC_MODULE_H__
 
-#include <map>
-#include <vector>
+#include <mutex>
 #include <string>
+#include <unordered_map>
 #include "workflow/WFTask.h" 
 #include "workflow/WFTaskFactory.h"
+#include "rpc_filter.h"
 
 namespace srpc
 {
 
-using RPCModuleData = std::map<std::string, std::string>;
-using RPCLogVector = std::vector<std::pair<std::string, std::string>>;
-
-static RPCModuleData global_empty_map;
-
 const char *const SRPC_SPAN_LOG			= "log";
 const char *const SRPC_SPAN_EVENT		= "event";
 const char *const SRPC_SPAN_MESSAGE		= "message";
 
-template<class REQ, class RESP>
 class RPCModule
 {
+protected:
+	virtual bool client_begin(SubTask *task, const RPCModuleData& data) = 0;
+	virtual bool server_begin(SubTask *task, const RPCModuleData& data) = 0;
+	virtual bool client_end(SubTask *task, const RPCModuleData& data) = 0;
+	virtual bool server_end(SubTask *task, const RPCModuleData& data) = 0;
+
 public:
-	virtual int client_begin(SubTask *task, const RPCModuleData& data) = 0;
-	virtual int server_begin(SubTask *task, const RPCModuleData& data) = 0;
+	bool client_task_begin(SubTask *task, const RPCModuleData& data)
+	{
+		bool ret = this->client_begin(task, data);
+
+		for (const auto& filter_pair : this->filters)
+			ret &= filter_pair.second->client_begin(task, data);
+
+		return ret;
+	}
+
+	bool server_task_begin(SubTask *task, const RPCModuleData& data)
+	{
+		bool ret = this->server_begin(task, data);
+
+		for (const auto& filter_pair : this->filters)
+			ret &= filter_pair.second->server_begin(task, data);
+
+		return ret;
+	}
+
+	bool client_task_end(SubTask *task, const RPCModuleData& data)
+	{
+		SubTask *filter_task;
+		bool ret = this->client_end(task, data);
+
+		for (const auto& filter_pair : this->filters)
+		{
+			ret &= filter_pair.second->client_end(task, data);
+			filter_task = filter_pair.second->create_filter_task(data);
+			series_of(task)->push_front(filter_task);
+		}
+
+		return ret;
+	}
 
-	virtual int client_end(SubTask *task, const RPCModuleData& data)
+	bool server_task_end(SubTask *task, const RPCModuleData& data)
 	{
-		SubTask *module_task = this->create_module_task(data);
-		series_of(task)->push_front(module_task);
-		return 0;
+		SubTask *filter_task;
+		bool ret = this->server_end(task, data);
+
+		for (const auto& filter_pair : this->filters)
+		{
+			ret &= filter_pair.second->server_end(task, data);
+			filter_task = filter_pair.second->create_filter_task(data);
+			series_of(task)->push_front(filter_task);
+		}
+
+		return ret;
 	}
 
-	virtual int server_end(SubTask *task, const RPCModuleData& data)
+	bool add_filter(RPCFilter *filter)
 	{
-		SubTask *module_task = this->create_module_task(data);
-		series_of(task)->push_front(module_task);
-		return 0;
+		this->mutex.lock();
+		auto iter = this->filters.find(filter->get_name());
+
+		if (iter == this->filters.end())
+			this->filters.insert(std::make_pair(filter->get_name(), filter));
+
+		this->mutex.unlock();
+		return iter == this->filters.end();
 	}
 
-	virtual SubTask* create_module_task(const RPCModuleData& data)
+	bool remove_filter(std::string name)
 	{
-		return WFTaskFactory::create_empty_task();
+		this->mutex.lock();
+		auto iter = this->filters.find(name);
+
+		if (iter != this->filters.end())
+			this->filters.erase(iter);
+
+		this->mutex.unlock();
+		return iter != this->filters.end();
 	}
+
+	size_t get_filters_size() const { return this->filters.size(); }
+	RPCModuleType get_module_type() const { return this->module_type; }
+	RPCModule(RPCModuleType module_type) { this->module_type = module_type; }
+
+private:
+	RPCModuleType module_type;
+	std::mutex mutex;
+	std::unordered_map<std::string, RPCFilter *> filters;
 };
 
 class SnowFlake

src/module/rpc_module_span.h

@@ -77,8 +77,7 @@ class RPCSpanContext : public RPCContextImpl<RPCREQ, RPCRESP>
 };
 
 template<class RPCTYPE>
-class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
-									   typename RPCTYPE::RESP>
+class RPCSpanModule : public RPCModule
 {
 public:
 	using CLIENT_TASK = RPCClientTask<typename RPCTYPE::REQ,
@@ -88,7 +87,7 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 	using SPAN_CONTEXT = RPCSpanContext<typename RPCTYPE::REQ,
 										typename RPCTYPE::RESP>;
 
-	int client_begin(SubTask *task, const RPCModuleData& data) override
+	bool client_begin(SubTask *task, const RPCModuleData& data) override
 	{
 		auto *client_task = static_cast<CLIENT_TASK *>(task);
 		auto *req = client_task->get_req();
@@ -115,13 +114,13 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 		module_data[SRPC_METHOD_NAME] = req->get_method_name();
 		module_data[SRPC_DATA_TYPE] = std::to_string(req->get_data_type());
 		module_data[SRPC_COMPRESS_TYPE] =
-										std::to_string(req->get_compress_type());
+								std::to_string(req->get_compress_type());
 		module_data[SRPC_START_TIMESTAMP] = std::to_string(GET_CURRENT_MS);
 
-		return 0; // always success
+		return true; // always success
 	}
 
-	int client_end(SubTask *task, const RPCModuleData& data) override
+	bool client_end(SubTask *task, const RPCModuleData& data) override
 	{
 		std::string ip;
 		unsigned short port;
@@ -141,13 +140,10 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 			module_data[SRPC_REMOTE_PORT] = std::to_string(port);
 		}
 
-		SubTask *module_task = this->create_module_task(module_data);
-		series_of(task)->push_front(module_task);
-
-		return 0;
+		return true;
 	}
 
-	int server_begin(SubTask *task, const RPCModuleData& data) override
+	bool server_begin(SubTask *task, const RPCModuleData& data) override
 	{
 		std::string ip;
 		unsigned short port;
@@ -159,7 +155,7 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 		module_data[SRPC_METHOD_NAME] = req->get_method_name();
 		module_data[SRPC_DATA_TYPE] = std::to_string(req->get_data_type());
 		module_data[SRPC_COMPRESS_TYPE] =
-										std::to_string(req->get_compress_type());
+								std::to_string(req->get_compress_type());
 
 		if (module_data[SRPC_TRACE_ID].empty())
 			module_data[SRPC_TRACE_ID] = std::to_string(
@@ -181,10 +177,10 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 		delete server_task->worker.ctx;
 		server_task->worker.ctx = span_ctx;
 
-		return 0;
+		return true;
 	}
 
-	int server_end(SubTask *task, const RPCModuleData& data) override
+	bool server_end(SubTask *task, const RPCModuleData& data) override
 	{
 		auto *server_task = static_cast<SERVER_TASK *>(task);
 		auto *resp = server_task->get_resp();
@@ -197,13 +193,64 @@ class RPCSpanModule : public RPCModule<typename RPCTYPE::REQ,
 		module_data[SRPC_STATE] = std::to_string(resp->get_status_code());
 		module_data[SRPC_ERROR] = std::to_string(resp->get_error());
 
-		SubTask *module_task = this->create_module_task(module_data);
-		series_of(task)->push_front(module_task);
+		return true;
+	}
+
+public:
+	RPCSpanModule() : RPCModule(RPCModuleSpan) { }
+};
 
-		return 0;
+template<class RPCTYPE>
+class RPCMonitorModule : public RPCModule
+{
+public:
+	bool client_begin(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
 	}
+	bool client_end(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+	bool server_begin(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+	bool server_end(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+
+public:
+	RPCMonitorModule() : RPCModule(RPCModuleMonitor) { }
+};
+
+template<class RPCTYPE>
+class RPCEmptyModule : public RPCModule
+{
+public:
+	bool client_begin(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+	bool client_end(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+	bool server_begin(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+	bool server_end(SubTask *task, const RPCModuleData& data) override
+	{
+		return true;
+	}
+
+public:
+	RPCEmptyModule() : RPCModule(RPCModuleMonitor) { }
 };
 
 } // end namespace srpc
 
 #endif
+