From 754bfb033a26a806eab204e6c9be8dcf87049382 Mon Sep 17 00:00:00 2001
From: Uku Loskit <ukuloskit@gmail.com>
Date: Sat, 30 Mar 2024 15:07:32 +0200
Subject: [PATCH] Problem: in some circumstances managed postgres data is not
 cleaned up

---
 extensions/omni_httpd/http_worker.c |  2 +-
 pg_yregress/instance.c              | 25 +++++++++++++++----------
 pg_yregress/pg_yregress.c           | 29 +++++++++++++++++++++++++++++
 pg_yregress/pg_yregress.h           |  1 +
 4 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/extensions/omni_httpd/http_worker.c b/extensions/omni_httpd/http_worker.c
index 037728338..176ea0d7b 100644
--- a/extensions/omni_httpd/http_worker.c
+++ b/extensions/omni_httpd/http_worker.c
@@ -280,7 +280,7 @@ void http_worker(Datum db_oid) {
           h2o_context_dispose(&iter.ref->context);
           MemoryContextDelete(iter.ref->memory_context);
           iter = clist_listener_contexts_erase_at(&listener_contexts, iter);
-        next_ctx: {}
+        next_ctx : {}
         }
       }
 
diff --git a/pg_yregress/instance.c b/pg_yregress/instance.c
index 06193fac6..b328f5ecd 100644
--- a/pg_yregress/instance.c
+++ b/pg_yregress/instance.c
@@ -204,14 +204,25 @@ yinstance_connect_result yinstance_connect(yinstance *instance) {
   return yinstance_connect_failure;
 }
 
-void yinstance_start(yinstance *instance) {
+void yinstance_prepare(yinstance *instance) {
   if (instance->managed) {
+    char datadir_template[] = "pg_yregress_XXXXXX";
+    char *datadir = mkdtemp(datadir_template);
+    if (datadir == NULL) {
+      fprintf(stderr, "Failed to create temporary directory: %s\n", strerror(errno));
+      return;
+    }
+    char *heap_datadir = strdup(datadir);
+    instance->info.managed.datadir = (iovec_t){.base = heap_datadir, .len = strlen(heap_datadir)};
+  }
+}
 
-    char datadir[] = "pg_yregress_XXXXXX";
-    mkdtemp(datadir);
+void yinstance_start(yinstance *instance) {
+  if (instance->managed) {
 
     // Initialize the cluster
     char *initdb_command;
+    const char *datadir = instance->info.managed.datadir.base;
     asprintf(&initdb_command,
              "%s/pg_ctl initdb -o '-A trust -U yregress --no-clean --no-sync --encoding=%.*s "
              "--locale=%.*s' -s -D %s",
@@ -278,11 +289,6 @@ void yinstance_start(yinstance *instance) {
     asprintf(&createdb_command, "%s/createdb -U yregress -O yregress -p %d yregress", bindir,
              instance->info.managed.port);
     system(createdb_command);
-
-    // Important to capture datadir before we try to start as init steps may restart the instance
-    // and it'll need the path
-    char *heap_datadir = strdup(datadir);
-    instance->info.managed.datadir = (iovec_t){.base = heap_datadir, .len = strlen(heap_datadir)};
   }
 
   // Wait until it is ready
@@ -316,6 +322,7 @@ static int remove_entry(const char *path, const struct stat *sb, int typeflag, s
 }
 
 void instances_cleanup() {
+  // TODO: clean up semaphores
   if (instances != NULL) {
     void *iter = NULL;
     struct fy_node_pair *instance_pair;
@@ -341,8 +348,6 @@ void instances_cleanup() {
   }
 }
 
-void register_instances_cleanup() { atexit(instances_cleanup); }
-
 void restart_instance(yinstance *instance) {
   if (instance->managed) {
     char *restart_command;
diff --git a/pg_yregress/pg_yregress.c b/pg_yregress/pg_yregress.c
index af1933322..9868e0922 100644
--- a/pg_yregress/pg_yregress.c
+++ b/pg_yregress/pg_yregress.c
@@ -1,6 +1,7 @@
 #include <getopt.h>
 #include <pwd.h>
 #include <stdio.h>
+#include <sys/wait.h>
 
 #include "pg_yregress.h"
 
@@ -521,6 +522,34 @@ static int execute_document(struct fy_document *fyd, bool managed, char *host, i
     }
   }
 
+  {
+    void *iter = NULL;
+    // prepare used instances, so we can actually start them in a child process
+    // the parent will handle clean up
+    struct fy_node_pair *instance_pair;
+    while ((instance_pair = fy_node_mapping_iterate(instances, &iter)) != NULL) {
+      yinstance *y_instance = (yinstance *)fy_node_get_meta(fy_node_pair_value(instance_pair));
+      if (y_instance->used) {
+        yinstance_prepare(y_instance);
+      }
+    }
+  }
+  pid_t pid = fork();
+  if (pid < 0) {
+    // TODO: still need to clean up empty datadir(s) that we have created
+    exit(EXIT_FAILURE);
+  }
+  if (pid > 0) {
+    // Parent process
+    int status;
+    waitpid(pid, &status, 0);
+    instances_cleanup();
+    return WEXITSTATUS(status);
+  }
+  // let's become the session leader, so if pg_yregress crashes, all the children would be
+  // terminated and the parent takes care of cleaning up
+  setsid();
+
   // Initialize used instances
   {
     fprintf(tap_file, "# Initializing instances\n");
diff --git a/pg_yregress/pg_yregress.h b/pg_yregress/pg_yregress.h
index 6102a0640..657217f2b 100644
--- a/pg_yregress/pg_yregress.h
+++ b/pg_yregress/pg_yregress.h
@@ -56,6 +56,7 @@ typedef enum {
   default_instance_ambiguous = 2
 } default_yinstance_result;
 
+void yinstance_prepare(yinstance *instance);
 void yinstance_start(yinstance *instance);
 
 typedef enum {