Make workerCount configurable for regression tests (#6764)

Make worker count flexible in our regression tests instead of hardcoding
it to 2 workers.

src/test/regress/Makefile

@@ -117,29 +117,31 @@ check-minimal-mx: all
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/mx_minimal_schedule $(EXTRA_TESTS)
 
 check-custom-schedule: all
-	$(pg_regress_multi_check) --load-extension=citus \
+	$(pg_regress_multi_check) --load-extension=citus --worker-count=$(WORKERCOUNT) \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
 check-failure-custom-schedule: all
-	$(pg_regress_multi_check) --load-extension=citus --mitmproxy \
+	$(pg_regress_multi_check) --load-extension=citus --mitmproxy --worker-count=$(WORKERCOUNT) \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
 check-isolation-custom-schedule: all  $(isolation_test_files)
-	$(pg_regress_multi_check) --load-extension=citus --isolationtester \
+	$(pg_regress_multi_check) --load-extension=citus --isolationtester --worker-count=$(WORKERCOUNT) \
 	-- $(MULTI_REGRESS_OPTS) --inputdir=$(citus_abs_srcdir)/build --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
 check-custom-schedule-vg: all
 	$(pg_regress_multi_check) --load-extension=citus \
-	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
+	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --worker-count=$(WORKERCOUNT) \
+	--valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
 check-failure-custom-schedule-vg: all
 	$(pg_regress_multi_check) --load-extension=citus --mitmproxy \
-	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
+	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --worker-count=$(WORKERCOUNT) \
+	--valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
 check-isolation-custom-schedule-vg: all  $(isolation_test_files)
-	$(pg_regress_multi_check) --load-extension=citus --isolationtester \
+	$(pg_regress_multi_check) --load-extension=citus --isolationtester --worker-count=$(WORKERCOUNT) \
 	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
 	-- $(MULTI_REGRESS_OPTS) --inputdir=$(citus_abs_srcdir)/build --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 

src/test/regress/citus_tests/run_test.py

@@ -73,10 +73,11 @@ if __name__ == "__main__":
         schedule: Optional[str]
         direct_extra_tests: list[str]
 
-        def __init__(self, schedule, extra_tests=None, repeatable=True):
+        def __init__(self, schedule, extra_tests=None, repeatable=True, worker_count=2):
             self.schedule = schedule
             self.direct_extra_tests = extra_tests or []
             self.repeatable = repeatable
+            self.worker_count = worker_count
 
         def extra_tests(self):
             all_deps = OrderedDict()
@@ -180,6 +181,15 @@ if __name__ == "__main__":
             return "base_schedule"
         return "minimal_schedule"
 
+    # we run the tests with 2 workers by default.
+    # If we find any dependency which requires more workers, we update the worker count.
+    def worker_count_for(test_name):
+        if test_name in deps:
+            return deps[test_name].worker_count
+        return 2
+
+    test_worker_count = max(worker_count_for(test_file_name), 2)
+
     if test_file_name in deps:
         dependencies = deps[test_file_name]
     elif schedule_line_is_upgrade_after(test_schedule_line):
@@ -204,6 +214,7 @@ if __name__ == "__main__":
     with open(tmp_schedule_path, "a") as myfile:
         for dependency in dependencies.extra_tests():
             myfile.write(f"test: {dependency}\n")
+            test_worker_count = max(worker_count_for(dependency), test_worker_count)
 
         repetition_cnt = args["repeat"]
         if repetition_cnt > 1 and not dependencies.repeatable:
@@ -224,7 +235,11 @@ if __name__ == "__main__":
         make_recipe += "-vg"
 
     # prepare command to run tests
-    test_command = f"make -C {regress_dir} {make_recipe} SCHEDULE='{pathlib.Path(tmp_schedule_path).stem}'"
+    test_command = (
+        f"make -C {regress_dir} {make_recipe} "
+        f"WORKERCOUNT={test_worker_count} "
+        f"SCHEDULE='{pathlib.Path(tmp_schedule_path).stem}'"
+    )
 
     # run test command n times
     try:

src/test/regress/pg_regress_multi.pl

@@ -49,6 +49,7 @@ sub Usage()
     print "  --pg_ctl-timeout    	Timeout for pg_ctl\n";
     print "  --connection-timeout	Timeout for connecting to worker nodes\n";
     print "  --mitmproxy        	Start a mitmproxy for one of the workers\n";
+    print "  --worker-count         Number of workers in Citus cluster (default: 2)\n";
     exit 1;
 }
 
@@ -84,6 +85,7 @@ my $mitmFifoPath = catfile($TMP_CHECKDIR, "mitmproxy.fifo");
 my $conninfo = "";
 my $publicWorker1Host = "localhost";
 my $publicWorker2Host = "localhost";
+my $workerCount = 2;
 
 my $serversAreShutdown = "TRUE";
 my $usingWindows = 0;
@@ -116,6 +118,7 @@ GetOptions(
     'conninfo=s' => \$conninfo,
     'worker-1-public-hostname=s' => \$publicWorker1Host,
     'worker-2-public-hostname=s' => \$publicWorker2Host,
+    'worker-count=i' => \$workerCount,
     'help' => sub { Usage() });
 
 my $fixopen = "$bindir/postgres.fixopen";
@@ -318,7 +321,6 @@ my $mitmPort = 9060;
 # Set some default configuration options
 my $masterPort = 57636;
 
-my $workerCount = 2;
 my @workerHosts = ();
 my @workerPorts = ();