1 files changed, 58 insertions, 31 deletions

diff --git a/daemon/pgsql.c b/daemon/pgsql.c
index 6bdb158..56eebd1 100644
--- a/daemon/pgsql.c
+++ b/daemon/pgsql.c
@@ -25,7 +25,7 @@ typedef struct pgsql_context
     /* Base Handler ------------------------------------------------------ */
     bd_context_t bd;
 
-    /* Settings ---------------------------------------------------------- */
+    /* Read Only --------------------------------------------------------- */
     const char* host;         /* The connection host or path */
     const char* port;         /* The connection port */
     const char* user;         /* The pgsql user name */
@@ -39,7 +39,7 @@ typedef struct pgsql_context
     int pgsql_max;            /* Number of open connections allowed */
     int pgsql_timeout;        /* Maximum amount of time to dedicate to a query */
 
-    /* Context ----------------------------------------------------------- */
+    /* Require Locking --------------------------------------------------- */
     PGconn** pool;            /* Pool of available connections */
     int pool_mark;          /* Amount of connections allocated */
 }
@@ -282,29 +282,42 @@ static const char* make_pgsql_connstring(const ha_request_t* rq, pgsql_context_t
 
 static PGconn* get_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx)
 {
-    PGconn* pg;
+    PGconn* pg = NULL;
     const char* connstring;
-    int i;
+    int i, create = 1;
 
     ASSERT(ctx);
 
-    for(i = 0; i < ctx->pgsql_max; i++)
-    {
-        /* An open connection in the pool */
-        if(ctx->pool[i])
+    /*
+     * Note that below there maybe a race condition between the two locks
+     * but this will only allow a few extra connections to open at best
+     * and as such really isn't a big issue.
+     */
+
+    ha_lock(NULL);
+
+        for(i = 0; i < ctx->pgsql_max; i++)
         {
-            ha_messagex(rq, LOG_DEBUG, "using cached pgsql connection");
-            pg = ctx->pool[i];
-            ctx->pool[i] = NULL;
-            return pg;
+            /* An open connection in the pool */
+            if(ctx->pool[i])
+            {
+                ha_messagex(rq, LOG_DEBUG, "using cached pgsql connection");
+                pg = ctx->pool[i];
+                ctx->pool[i] = NULL;
+                break;
+            }
         }
-    }
 
-    if(ctx->pool_mark >= ctx->pgsql_max)
-    {
-        ha_messagex(rq, LOG_ERR, "too many open pgsql connections");
-        return NULL;
-    }
+        if(pg == NULL && ctx->pool_mark >= ctx->pgsql_max)
+        {
+            ha_messagex(rq, LOG_ERR, "too many open pgsql connections");
+            create = 0;
+        }
+
+    ha_unlock(NULL);
+
+    if(pg != NULL || create == 0)
+        return pg;
 
     connstring = make_pgsql_connstring(rq, ctx);
     if(!connstring)
@@ -326,16 +339,26 @@ static PGconn* get_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx
         return NULL;
     }
 
-    ctx->pool_mark++;
-    ha_messagex(rq, LOG_DEBUG, "opened new pgsql connection (total %d)", ctx->pool_mark);
+    ha_lock(NULL);
+
+        ctx->pool_mark++;
+        ha_messagex(rq, LOG_DEBUG, "opened new pgsql connection (total %d)", ctx->pool_mark);
+
+    ha_unlock(NULL);
+
     return pg;
 }
 
 static void discard_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx, PGconn* pg)
 {
     PQfinish(pg);
-    ctx->pool_mark--;
-    ha_messagex(rq, LOG_DEBUG, "discarding pgsql connection (total %d)", ctx->pool_mark);
+
+    ha_lock(NULL);
+
+        ctx->pool_mark--;
+        ha_messagex(rq, LOG_DEBUG, "discarding pgsql connection (total %d)", ctx->pool_mark);
+
+    ha_unlock(NULL);
 }
 
 static void save_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx, PGconn* pg)
@@ -350,17 +373,21 @@ static void save_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx,
     /* Make sure it's worth saving */
     if(PQstatus(pg) != CONNECTION_BAD)
     {
-        for(i = 0; i < ctx->pgsql_max; i++)
-        {
-            /* An open connection in the pool */
-            if(!ctx->pool[i])
+        ha_lock(NULL);
+
+            for(i = 0; i < ctx->pgsql_max; i++)
             {
-                ha_messagex(rq, LOG_DEBUG, "caching pgsql connection for later use");
-                ctx->pool[i] = pg;
-                pg = NULL;
-                break;
+                /* An open connection in the pool */
+                if(!ctx->pool[i])
+                {
+                    ha_messagex(rq, LOG_DEBUG, "caching pgsql connection for later use");
+                    ctx->pool[i] = pg;
+                    pg = NULL;
+                    break;
+                }
             }
-        }
+
+        ha_unlock(NULL);
     }
 
     if(pg != NULL)