From 01a537305a7556c0f9d297019f8820b1366786e6 Mon Sep 17 00:00:00 2001
From: Yuri Astrakhan <YuriAstrakhan@gmail.com>
Date: Tue, 28 Apr 2020 16:20:28 -0400
Subject: [PATCH] [WIP] Limiting postserve PG connection count

This is a (currently non-working) support to limit
the maximum number of simultaneous PostgreSQL connection.
Probably due to how web serving framework functions,
more than allowed PG connections are still made.

Feedback / ideas are welcome.
---
 bin/postserve             |  3 +++
 openmaptiles/postserve.py | 15 ++++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/bin/postserve b/bin/postserve
index 4df26916..9b4e553e 100755
--- a/bin/postserve
+++ b/bin/postserve
@@ -9,6 +9,7 @@ Usage:
                       [--pghost=<host>] [--pgport=<port>] [--dbname=<db>]
                       [--user=<user>] [--password=<password>]
                       [--test-geometry] [--verbose]
+                      [--max-pg-conn <count>]
   postserve --help
   postserve --version
 
@@ -24,6 +25,7 @@ Options:
   --no-feature-ids      Disable feature ID generation, e.g. from osm_id.
                         Feature IDS are automatically disabled with PostGIS before v3
   -g --test-geometry    Validate all geometries produced by ST_AsMvtGeom(), and warn.
+  -c --max-pg-conn <c>  Maximum number of PostgreSQL connections. Defaults to CPU count.
   -v --verbose          Print additional debugging information
   --help                Show this screen.
   --version             Show version.
@@ -66,6 +68,7 @@ def main(args):
         disable_feature_ids=args['--no-feature-ids'],
         test_geometry=args['--test-geometry'],
         verbose=args.get('--verbose'),
+        max_pg_connections=args.get('--max-pg-conn'),
     ).serve()
 
 
diff --git a/openmaptiles/postserve.py b/openmaptiles/postserve.py
index fe41ab0c..d50ddd9d 100644
--- a/openmaptiles/postserve.py
+++ b/openmaptiles/postserve.py
@@ -1,4 +1,5 @@
 import logging
+import os
 from functools import partial
 from typing import Union, List, Any, Dict
 
@@ -137,7 +138,7 @@ class Postserve:
 
     def __init__(self, url, port, pghost, pgport, dbname, user, password,
                  layers, tileset_path, sql_file, key_column, disable_feature_ids,
-                 gzip, verbose, exclude_layers, test_geometry):
+                 gzip, verbose, exclude_layers, test_geometry, max_pg_connections):
         self.url = url
         self.port = port
         self.pghost = pghost
@@ -155,6 +156,10 @@ def __init__(self, url, port, pghost, pgport, dbname, user, password,
         self.test_geometry = test_geometry
         self.verbose = verbose
 
+        if max_pg_connections is None:
+            max_pg_connections = os.cpu_count() or 1
+        self.max_pg_connections = max_pg_connections
+
         self.tileset = Tileset.parse(self.tileset_path)
 
     def create_metadata(self,
@@ -254,12 +259,16 @@ def serve(self):
         access_log.setLevel(logging.INFO if self.verbose else logging.ERROR)
 
         print(f'Connecting to PostgreSQL at {self.pghost}:{self.pgport}, '
-              f'db={self.dbname}, user={self.user}...')
+              f'db={self.dbname}, user={self.user}, '
+              f'up to {self.max_pg_connections} simultaneous connections...')
         io_loop = IOLoop.current()
         self.pool = io_loop.run_sync(partial(
             create_pool,
             dsn=f"postgresql://{self.user}:{self.password}@"
-                f"{self.pghost}:{self.pgport}/{self.dbname}"))
+                f"{self.pghost}:{self.pgport}/{self.dbname}",
+            min_size=min(2, self.max_pg_connections),
+            max_size=self.max_pg_connections,
+        ))
         io_loop.run_sync(partial(self.init_connection))
 
         if self.sql_file: