From b1bd2103d763e6099464f2712676d411822fd0b2 Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Tue, 10 Oct 2023 14:12:39 +0300
Subject: [PATCH] opt(server): If unspecified, use num_shards ==
 proactor_threads (#1998)

* opt(server): If unspecified, use num_shards == proactor_threads

* Use same config for tests, per Roman's request
---
 helio                       | 2 +-
 src/server/main_service.cc  | 7 ++++---
 src/server/test_utils.cc    | 5 +++++
 tests/dragonfly/instance.py | 8 ++++++++
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/helio b/helio
index 6655f713c..0950c8fd6 160000
--- a/helio
+++ b/helio
@@ -1 +1 @@
-Subproject commit 6655f713c7cb6aebc586973ebea6d6bfea06df3c
+Subproject commit 0950c8fd612bba442caffdabd86d67a32246d461
diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index 4b67f921d..5a42268f8 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -686,9 +686,10 @@ void Service::Init(util::AcceptServer* acceptor, std::vector<facade::Listener*>
   pp_.Await([reg](uint32_t index, ProactorBase* pb) { ServerState::Init(index, reg); });
 
   uint32_t shard_num = GetFlag(FLAGS_num_shards);
-  if (shard_num == 0) {
-    shard_num = pp_.size() > 1 ? pp_.size() - 1 : pp_.size();
-  } else if (shard_num > pp_.size()) {
+  if (shard_num == 0 || shard_num > pp_.size()) {
+    LOG_IF(WARNING, shard_num > pp_.size())
+        << "Requested num_shards (" << shard_num << ") is bigger than thread count (" << pp_.size()
+        << "), using num_shards=" << pp_.size();
     shard_num = pp_.size();
   }
 
diff --git a/src/server/test_utils.cc b/src/server/test_utils.cc
index 6c2d50c91..cc4d34100 100644
--- a/src/server/test_utils.cc
+++ b/src/server/test_utils.cc
@@ -25,6 +25,7 @@ extern "C" {
 using namespace std;
 
 ABSL_DECLARE_FLAG(string, dbfilename);
+ABSL_DECLARE_FLAG(uint32_t, num_shards);
 ABSL_FLAG(bool, force_epoll, false, "If true, uses epoll api instead iouring to run tests");
 
 namespace dfly {
@@ -193,6 +194,10 @@ void BaseFamilyTest::ResetService() {
   pp_.reset(fb2::Pool::Epoll(num_threads_));
 #endif
 
+  // Using a different default than production could expose bugs
+  if (absl::GetFlag(FLAGS_num_shards) == 0) {
+    absl::SetFlag(&FLAGS_num_shards, num_threads_ - 1);
+  }
   pp_->Run();
   service_ = std::make_unique<Service>(pp_.get());
 
diff --git a/tests/dragonfly/instance.py b/tests/dragonfly/instance.py
index 8b8ab1eeb..20a89c462 100644
--- a/tests/dragonfly/instance.py
+++ b/tests/dragonfly/instance.py
@@ -74,6 +74,14 @@ class DflyInstance:
             del self.args["logtostderr"]
             self.args["alsologtostderr"] = None
 
+        # Run with num_shards = (proactor_threads - 1) if possible, so help expose bugs
+        if "num_shards" not in self.args:
+            threads = psutil.cpu_count()
+            if "proactor_threads" in self.args:
+                threads = int(self.args["proactor_threads"])
+            if threads > 1:
+                self.args["num_shards"] = threads - 1
+
     def __del__(self):
         assert self.proc == None