[#110] Assign pid before the bad connection event can happen

agroal · Oct 12, 2020 · 39a1017 · 39a1017
1 parent 6692f0f
commit 39a1017
Showing 1 changed file with 13 additions and 17 deletions.
diff --git a/src/libpgagroal/pool.c b/src/libpgagroal/pool.c
@@ -61,7 +61,6 @@ static bool do_prefill(void* shmem, char* username, char* database, int size);
 int
 pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse, bool transaction_mode, int* slot)
 {
-   bool prefill;
    bool do_init;
    bool has_lock;
    int connections;
@@ -76,8 +75,6 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
 
    config = (struct configuration*)shmem;
 
-   prefill = false;
-
    pgagroal_prometheus_connection_get(shmem);
 
    best_rule = find_best_rule(shmem, username, database);
@@ -146,11 +143,16 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
 
    if (*slot != -1)
    {
+      config->connections[*slot].limit_rule = best_rule;
+      config->connections[*slot].pid = getpid();
+
       if (do_init)
       {
          /* We need to find the server for the connection */
          if (pgagroal_get_primary(shmem, &server))
          {
+            config->connections[*slot].limit_rule = -1;
+            config->connections[*slot].pid = -1;
             atomic_store(&config->states[*slot], STATE_NOTINIT);
 
             if (!fork())
@@ -166,6 +168,8 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
          if (pgagroal_connect(shmem, config->servers[server].host, config->servers[server].port, &fd))
          {
             ZF_LOGE("pgagroal: No connection to %s:%d", config->servers[server].host, config->servers[server].port);
+            config->connections[*slot].limit_rule = -1;
+            config->connections[*slot].pid = -1;
             atomic_store(&config->states[*slot], STATE_NOTINIT);
 
             pgagroal_prometheus_server_error(server, shmem);
@@ -195,10 +199,7 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
          memset(&config->connections[*slot].database, 0, MAX_DATABASE_LENGTH);
          memcpy(&config->connections[*slot].database, database, MIN(strlen(database), MAX_DATABASE_LENGTH - 1));
 
-         config->connections[*slot].limit_rule = best_rule;
          config->connections[*slot].has_security = SECURITY_INVALID;
-         config->connections[*slot].timestamp = time(NULL);
-         config->connections[*slot].pid = getpid();
          config->connections[*slot].fd = fd;
 
          atomic_store(&config->states[*slot], STATE_IN_USE);
@@ -232,7 +233,12 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
 
             ZF_LOGD("pgagroal_get_connection: Slot %d FD %d - Error", *slot, config->connections[*slot].fd);
             status = pgagroal_kill_connection(shmem, *slot);
-            prefill = true;
+
+            if (!fork())
+            {
+               pgagroal_prefill(shmem, false);
+            }
+
             if (status == 0)
             {
                goto retry2;
@@ -244,16 +250,6 @@ pgagroal_get_connection(void* shmem, char* username, char* database, bool reuse,
          }
       }
 
-      if (prefill)
-      {
-         if (!fork())
-         {
-            pgagroal_prefill(shmem, false);
-         }
-      }
-
-      config->connections[*slot].limit_rule = best_rule;
-      config->connections[*slot].pid = getpid();
       config->connections[*slot].timestamp = time(NULL);
 
       pgagroal_prometheus_connection_success(shmem);