Merge pull request #874 from Barenboim/master

add 'pre_select' for WFServiceGovernance

src/kernel/thrdpool.c

@@ -207,7 +207,10 @@ thrdpool_t *thrdpool_create(size_t nthreads, size_t stacksize)
 }
 
 inline void __thrdpool_schedule(const struct thrdpool_task *task, void *buf,
-								thrdpool_t *pool)
+								thrdpool_t *pool);
+
+void __thrdpool_schedule(const struct thrdpool_task *task, void *buf,
+						 thrdpool_t *pool)
 {
 	struct __thrdpool_task_entry *entry = (struct __thrdpool_task_entry *)buf;
 

src/nameservice/WFDnsResolver.cc

@@ -362,7 +362,7 @@ void WFResolverTask::dispatch()
 		series_of(this)->push_front(dns_task);
 	}
 
-	query_dns_ = true;
+	has_next_ = true;
 	this->subtask_done();
 }
 
@@ -370,13 +370,15 @@ SubTask *WFResolverTask::done()
 {
 	SeriesWork *series = series_of(this);
 
-	if (!query_dns_)
+	if (!has_next_)
 	{
 		if (this->callback)
 			this->callback(this);
 
 		delete this;
 	}
+	else
+		has_next_ = false;
 
 	return series->pop();
 }

src/nameservice/WFDnsResolver.h

@@ -37,7 +37,7 @@ public:
 	{
 		dns_ttl_default_ = dns_ttl_default;
 		dns_ttl_min_ = dns_ttl_min;
-		query_dns_ = false;
+		has_next_ = false;
 	}
 
 	WFResolverTask(const struct WFNSParams *ns_params,
@@ -45,12 +45,13 @@ public:
 		WFRouterTask(std::move(cb)),
 		ns_params_(*ns_params)
 	{
-		query_dns_ = false;
+		has_next_ = false;
 	}
 
 protected:
 	virtual void dispatch();
 	virtual SubTask *done();
+	void set_has_next() { has_next_ = true; }
 
 private:
 	void thread_dns_callback(void *thrd_dns_task);
@@ -70,7 +71,7 @@ protected:
 private:
 	const char *host_;
 	unsigned short port_;
-	bool query_dns_;
+	bool has_next_;
 };
 
 class WFDnsResolver : public WFNSPolicy

src/nameservice/WFNameService.h

@@ -35,6 +35,10 @@ class WFRouterTask : public WFGenericTask
 public:
 	RouteManager::RouteResult *get_result() { return &this->result; }
 
+public:
+	void set_state(int state) { this->state = state; }
+	void set_error(int error) { this->error = error; }
+
 protected:
 	RouteManager::RouteResult result;
 	std::function<void (WFRouterTask *)> callback;

src/nameservice/WFServiceGovernance.cc

@@ -116,6 +116,23 @@ void WFSGResolverTask::dispatch()
 	WFNSTracing *tracing = ns_params_.tracing;
 	EndpointAddress *addr;
 
+	if (sg_->pre_select_)
+	{
+		WFConditional *cond = sg_->pre_select_(this);
+		if (cond)
+		{
+			series_of(this)->push_front(cond);
+			this->set_has_next();
+			this->subtask_done();
+			return;
+		}
+		else if (this->state != WFT_STATE_UNDEFINED)
+		{
+			this->subtask_done();
+			return;
+		}
+	}
+
 	if (sg_->select(ns_params_.uri, tracing, &addr))
 	{
 		auto *tracing_data = (WFServiceGovernance::TracingData *)tracing->data;
@@ -134,13 +151,12 @@ void WFSGResolverTask::dispatch()
 		dns_ttl_min_ = addr->params->dns_ttl_min;
 		ep_params_ = addr->params->endpoint_params;
 		this->WFResolverTask::dispatch();
-		return;
 	}
 	else
 	{
 		this->state = WFT_STATE_TASK_ERROR;
 		this->error = WFT_ERR_UPSTREAM_UNAVAILABLE;
-		return this->subtask_done();
+		this->subtask_done();
 	}
 }
 

src/nameservice/WFServiceGovernance.h

@@ -25,6 +25,7 @@
 #include <unordered_map>
 #include <vector>
 #include <atomic>
+#include <functional>
 #include "URIParser.h"
 #include "EndpointParams.h"
 #include "WFNameService.h"
@@ -128,6 +129,14 @@ public:
 	void set_mttr_second(unsigned int second) { this->mttr_second = second; }
 	static bool in_select_history(WFNSTracing *tracing, EndpointAddress *addr);
 
+public:
+	using pre_select_t = std::function<WFConditional *(WFRouterTask *)>;
+
+	void set_pre_select(pre_select_t pre_select)
+	{
+		pre_select_ = std::move(pre_select);
+	}
+
 public:
 	WFServiceGovernance() :
 		breaker_lock(PTHREAD_MUTEX_INITIALIZER),
@@ -166,9 +175,11 @@ private:
 	void fuse_server_to_breaker(EndpointAddress *addr);
 	void check_breaker_locked(int64_t cur_time);
 
+private:
 	struct list_head breaker_list;
 	pthread_mutex_t breaker_lock;
 	unsigned int mttr_second;
+	pre_select_t pre_select_;
 
 protected:
 	virtual EndpointAddress *first_strategy(const ParsedURI& uri,