From b6eaff518afc7708436406d76335a4f2031eac45 Mon Sep 17 00:00:00 2001
From: XieHan <XieHan@sogou-inc.com>
Date: Sat, 19 Jun 2021 01:09:14 +0800
Subject: [PATCH] update CommMessageIn::feedback. Update WFMailboxTask

---
 src/factory/WFTask.h        | 51 +++++++++++++++++++++++++------------
 src/factory/WFTaskFactory.h |  6 +++++
 src/kernel/Communicator.cc  |  2 +-
 src/kernel/Communicator.h   |  2 +-
 4 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/src/factory/WFTask.h b/src/factory/WFTask.h
index 0369c310..421fa946 100644
--- a/src/factory/WFTask.h
+++ b/src/factory/WFTask.h
@@ -556,20 +556,18 @@ protected:
 	virtual ~WFCounterTask() { }
 };
 
-class WFMailboxTask : public WFCounterTask
+class WFMailboxTask : public WFGenericTask
 {
 public:
-	void send(void *mail)
+	void send(void *msg)
 	{
-		*this->next++ = mail;
+		*this->next++ = msg;
 		this->count();
 	}
 
-	void *get_mail(size_t index) { return this->mailbox[index]; }
-
-	void **get_mailbox(size_t *count)
+	void **get_mailbox(size_t *n)
 	{
-		*count = this->next - this->mailbox;
+		*n = this->next - this->mailbox;
 		return this->mailbox;
 	}
 
@@ -579,23 +577,44 @@ public:
 		this->callback = std::move(cb);
 	}
 
+public:
+	virtual void count()
+	{
+		if (--this->value == 0)
+		{
+			this->state = WFT_STATE_SUCCESS;
+			this->subtask_done();
+		}
+	}
+
+protected:
+	virtual void dispatch()
+	{
+		this->WFMailboxTask::count();
+	}
+
+	virtual SubTask *done()
+	{
+		SeriesWork *series = series_of(this);
+
+		if (this->callback)
+			this->callback(this);
+
+		delete this;
+		return series->pop();
+	}
+
 protected:
 	void **mailbox;
 	std::atomic<void **> next;
+	std::atomic<size_t> value;
 	std::function<void (WFMailboxTask *)> callback;
 
-private:
-	static void callback_wrapper(WFCounterTask *task)
-	{
-		WFMailboxTask *mbt = (WFMailboxTask *)task;
-		mbt->callback(mbt);
-	}
-
 public:
-	WFMailboxTask(void *mailbox[], size_t size,
+	WFMailboxTask(void **mailbox, size_t size,
 				  std::function<void (WFMailboxTask *)>&& cb) :
-		WFCounterTask(size, WFMailboxTask::callback_wrapper),
 		next(mailbox),
+		value(size + 1),
 		callback(std::move(cb))
 	{
 		this->mailbox = mailbox;
diff --git a/src/factory/WFTaskFactory.h b/src/factory/WFTaskFactory.h
index 39042a11..cc88f2e7 100644
--- a/src/factory/WFTaskFactory.h
+++ b/src/factory/WFTaskFactory.h
@@ -90,6 +90,8 @@ using fsync_callback_t = std::function<void (WFFileSyncTask *)>;
 using timer_callback_t = std::function<void (WFTimerTask *)>;
 using counter_callback_t = std::function<void (WFCounterTask *)>;
 
+using mailbox_callback_t = std::function<void (WFMailboxTask *)>;
+
 // Graph (DAG) task.
 using graph_callback_t = std::function<void (WFGraphTask *)>;
 
@@ -223,6 +225,10 @@ public:
 	 * creation, and more than one counter may reach target value. */
 	static void count_by_name(const std::string& counter_name, unsigned int n);
 
+public:
+	static WFMailboxTask *create_mailbox_task(size_t size,
+											  mailbox_callback_t callback);
+
 public:
 	template<class FUNC, class... ARGS>
 	static WFGoTask *create_go_task(const std::string& queue_name,
diff --git a/src/kernel/Communicator.cc b/src/kernel/Communicator.cc
index 3b266642..8770bd74 100644
--- a/src/kernel/Communicator.cc
+++ b/src/kernel/Communicator.cc
@@ -154,7 +154,7 @@ void CommTarget::deinit()
 	free(this->addr);
 }
 
-int CommMessageIn::feedback(const char *buf, size_t size)
+int CommMessageIn::feedback(const void *buf, size_t size)
 {
 	struct CommConnEntry *entry = this->entry;
 	int error;
diff --git a/src/kernel/Communicator.h b/src/kernel/Communicator.h
index 7a617520..fa12597b 100644
--- a/src/kernel/Communicator.h
+++ b/src/kernel/Communicator.h
@@ -110,7 +110,7 @@ private:
 
 protected:
 	/* Send small packet while receiving. Call only in append(). */
-	int feedback(const char *buf, size_t size);
+	int feedback(const void *buf, size_t size);
 
 private:
 	struct CommConnEntry *entry;