perf: remove repeated GetSession calls for FixedSizePool

Add a _last_use_time to Session and use this to determine whether the
FixedSizePool should check whether the session still exists, and whether
it should be replaced. This significantly reduces the number of times that
GetSession is called when using FixedSizePool.

Author: olavloite

google/cloud/spanner_v1/pool.py

@@ -26,7 +26,12 @@
 )
 from warnings import warn
 
-_NOW = datetime.datetime.utcnow  # unit tests may replace
+
+def _now():
+    datetime.datetime.now(datetime.UTC)
+
+
+_NOW = _now  # unit tests may replace
 
 
 class AbstractSessionPool(object):
@@ -145,7 +150,8 @@ class FixedSizePool(AbstractSessionPool):
     - Pre-allocates / creates a fixed number of sessions.
 
     - "Pings" existing sessions via :meth:`session.exists` before returning
-      them, and replaces expired sessions.
+      sessions that have not been used for more than 55 minutes and replaces
+      expired sessions.
 
     - Blocks, with a timeout, when :meth:`get` is called on an empty pool.
       Raises after timing out.
@@ -171,18 +177,21 @@ class FixedSizePool(AbstractSessionPool):
 
     DEFAULT_SIZE = 10
     DEFAULT_TIMEOUT = 10
+    DEFAULT_MAX_AGE_MINUTES = 55
 
     def __init__(
         self,
         size=DEFAULT_SIZE,
         default_timeout=DEFAULT_TIMEOUT,
         labels=None,
         database_role=None,
+        max_age_minutes=DEFAULT_MAX_AGE_MINUTES,
     ):
         super(FixedSizePool, self).__init__(labels=labels, database_role=database_role)
         self.size = size
         self.default_timeout = default_timeout
         self._sessions = queue.LifoQueue(size)
+        self._max_age = datetime.timedelta(minutes=max_age_minutes)
 
     def bind(self, database):
         """Associate the pool with a database.
@@ -230,8 +239,9 @@ def get(self, timeout=None):
             timeout = self.default_timeout
 
         session = self._sessions.get(block=True, timeout=timeout)
+        age = _NOW() - session.last_use_time
 
-        if not session.exists():
+        if age >= self._max_age and not session.exists():
             session = self._database.session()
             session.create()
 

google/cloud/spanner_v1/session.py

@@ -17,6 +17,7 @@
 from functools import total_ordering
 import random
 import time
+from datetime import datetime, UTC
 
 from google.api_core.exceptions import Aborted
 from google.api_core.exceptions import GoogleAPICallError
@@ -69,6 +70,7 @@ def __init__(self, database, labels=None, database_role=None):
             labels = {}
         self._labels = labels
         self._database_role = database_role
+        self._last_use_time = datetime.now(UTC)
 
     def __lt__(self, other):
         return self._session_id < other._session_id
@@ -78,6 +80,14 @@ def session_id(self):
         """Read-only ID, set by the back-end during :meth:`create`."""
         return self._session_id
 
+    @property
+    def last_use_time(self):
+        """ "Approximate last use time of this session
+
+        :rtype: datetime
+        :returns: the approximate last use time of this session"""
+        return self._last_use_time
+
     @property
     def database_role(self):
         """User-assigned database-role for the session.
@@ -154,6 +164,7 @@ def create(self):
                 metadata=metadata,
             )
         self._session_id = session_pb.name.split("/")[-1]
+        self._last_use_time = datetime.now()
 
     def exists(self):
         """Test for the existence of this session.
@@ -181,6 +192,7 @@ def exists(self):
         ) as span:
             try:
                 api.get_session(name=self.name, metadata=metadata)
+                self._last_use_time = datetime.now()
                 if span:
                     span.set_attribute("session_found", True)
             except NotFound:
@@ -222,6 +234,7 @@ def ping(self):
         metadata = _metadata_with_prefix(self._database.name)
         request = ExecuteSqlRequest(session=self.name, sql="SELECT 1")
         api.execute_sql(request=request, metadata=metadata)
+        self._last_use_time = datetime.now()
 
     def snapshot(self, **kw):
         """Create a snapshot to perform a set of reads with shared staleness.
@@ -273,6 +286,7 @@ def read(self, table, columns, keyset, index="", limit=0, column_info=None):
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
+        self._last_use_time = datetime.now()
         return self.snapshot().read(
             table, columns, keyset, index, limit, column_info=column_info
         )
@@ -339,6 +353,7 @@ def execute_sql(
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
+        self._last_use_time = datetime.now()
         return self.snapshot().execute_sql(
             sql,
             params,
@@ -378,6 +393,7 @@ def transaction(self):
             del self._transaction
 
         txn = self._transaction = Transaction(self)
+        self._last_use_time = datetime.now()
         return txn
 
     def run_in_transaction(self, func, *args, **kw):
@@ -444,6 +460,7 @@ def run_in_transaction(self, func, *args, **kw):
                 raise
 
             try:
+                self._last_use_time = datetime.now()
                 txn.commit(
                     return_commit_stats=self._database.log_commit_stats,
                     request_options=commit_request_options,

google/cloud/spanner_v1/testing/mock_database_admin.py

@@ -16,6 +16,7 @@
 from google.protobuf import empty_pb2
 import google.cloud.spanner_v1.testing.spanner_database_admin_pb2_grpc as database_admin_grpc
 
+
 # An in-memory mock DatabaseAdmin server that can be used for testing.
 class DatabaseAdminServicer(database_admin_grpc.DatabaseAdminServicer):
     def __init__(self):

google/cloud/spanner_v1/testing/mock_spanner.py

@@ -24,6 +24,7 @@
 import google.cloud.spanner_v1.types.spanner as spanner
 import google.cloud.spanner_v1.types.transaction as transaction
 
+
 class MockSpanner:
     def __init__(self):
         self.results = {}

tests/mockserver_tests/test_basics.py

@@ -29,7 +29,6 @@
     FixedSizePool,
     BatchCreateSessionsRequest,
     ExecuteSqlRequest,
-    GetSessionRequest,
 )
 from google.cloud.spanner_v1.database import Database
 from google.cloud.spanner_v1.instance import Instance
@@ -125,12 +124,9 @@ def test_select1(self):
                 self.assertEqual(1, row[0])
             self.assertEqual(1, len(result_list))
         requests = self.spanner_service.requests
-        self.assertEqual(3, len(requests))
+        self.assertEqual(2, len(requests), msg=requests)
         self.assertTrue(isinstance(requests[0], BatchCreateSessionsRequest))
-        # TODO: Optimize FixedSizePool so this GetSessionRequest is not executed
-        #       every time a session is fetched.
-        self.assertTrue(isinstance(requests[1], GetSessionRequest))
-        self.assertTrue(isinstance(requests[2], ExecuteSqlRequest))
+        self.assertTrue(isinstance(requests[1], ExecuteSqlRequest))
 
     def test_create_table(self):
         database_admin_api = self.client.database_admin_api

tests/unit/test_pool.py

@@ -15,6 +15,7 @@
 
 from functools import total_ordering
 import unittest
+from datetime import datetime, UTC, timedelta
 
 import mock
 
@@ -184,13 +185,30 @@ def test_bind(self):
         for session in SESSIONS:
             session.create.assert_not_called()
 
-    def test_get_non_expired(self):
+    def test_get_active(self):
         pool = self._make_one(size=4)
         database = _Database("name")
         SESSIONS = sorted([_Session(database) for i in range(0, 4)])
         database._sessions.extend(SESSIONS)
         pool.bind(database)
 
+        # check if sessions returned in LIFO order
+        for i in (3, 2, 1, 0):
+            session = pool.get()
+            self.assertIs(session, SESSIONS[i])
+            self.assertFalse(session._exists_checked)
+            self.assertFalse(pool._sessions.full())
+
+    def test_get_non_expired(self):
+        pool = self._make_one(size=4)
+        database = _Database("name")
+        last_use_time = datetime.now(UTC) - timedelta(minutes=56)
+        SESSIONS = sorted(
+            [_Session(database, last_use_time=last_use_time) for i in range(0, 4)]
+        )
+        database._sessions.extend(SESSIONS)
+        pool.bind(database)
+
         # check if sessions returned in LIFO order
         for i in (3, 2, 1, 0):
             session = pool.get()
@@ -201,7 +219,8 @@ def test_get_non_expired(self):
     def test_get_expired(self):
         pool = self._make_one(size=4)
         database = _Database("name")
-        SESSIONS = [_Session(database)] * 5
+        last_use_time = datetime.now(UTC) - timedelta(minutes=65)
+        SESSIONS = [_Session(database, last_use_time=last_use_time)] * 5
         SESSIONS[0]._exists = False
         database._sessions.extend(SESSIONS)
         pool.bind(database)
@@ -497,7 +516,7 @@ def test_get_hit_w_ping(self):
         SESSIONS = [_Session(database)] * 4
         database._sessions.extend(SESSIONS)
 
-        sessions_created = datetime.datetime.utcnow() - datetime.timedelta(seconds=4000)
+        sessions_created = datetime.datetime.now(UTC) - datetime.timedelta(seconds=4000)
 
         with _Monkey(MUT, _NOW=lambda: sessions_created):
             pool.bind(database)
@@ -519,7 +538,7 @@ def test_get_hit_w_ping_expired(self):
         SESSIONS[0]._exists = False
         database._sessions.extend(SESSIONS)
 
-        sessions_created = datetime.datetime.utcnow() - datetime.timedelta(seconds=4000)
+        sessions_created = datetime.datetime.now(UTC) - datetime.timedelta(seconds=4000)
 
         with _Monkey(MUT, _NOW=lambda: sessions_created):
             pool.bind(database)
@@ -575,7 +594,7 @@ def test_put_non_full(self):
         pool = self._make_one(size=1)
         session_queue = pool._sessions = _Queue()
 
-        now = datetime.datetime.utcnow()
+        now = datetime.datetime.now(UTC)
         database = _Database("name")
         session = _Session(database)
 
@@ -631,7 +650,7 @@ def test_ping_oldest_stale_but_exists(self):
         database._sessions.extend(SESSIONS)
         pool.bind(database)
 
-        later = datetime.datetime.utcnow() + datetime.timedelta(seconds=4000)
+        later = datetime.datetime.now(UTC) + datetime.timedelta(seconds=4000)
         with _Monkey(MUT, _NOW=lambda: later):
             pool.ping()
 
@@ -649,7 +668,7 @@ def test_ping_oldest_stale_and_not_exists(self):
         database._sessions.extend(SESSIONS)
         pool.bind(database)
 
-        later = datetime.datetime.utcnow() + datetime.timedelta(seconds=4000)
+        later = datetime.datetime.now(UTC) + datetime.timedelta(seconds=4000)
         with _Monkey(MUT, _NOW=lambda: later):
             pool.ping()
 
@@ -733,7 +752,7 @@ def test_bind_w_timestamp_race(self):
         from google.cloud._testing import _Monkey
         from google.cloud.spanner_v1 import pool as MUT
 
-        NOW = datetime.datetime.utcnow()
+        NOW = datetime.datetime.now(UTC)
         pool = self._make_one()
         database = _Database("name")
         SESSIONS = [_Session(database) for _ in range(10)]
@@ -915,18 +934,25 @@ def _make_transaction(*args, **kw):
 class _Session(object):
     _transaction = None
 
-    def __init__(self, database, exists=True, transaction=None):
+    def __init__(
+        self, database, exists=True, transaction=None, last_use_time=datetime.now(UTC)
+    ):
         self._database = database
         self._exists = exists
         self._exists_checked = False
         self._pinged = False
         self.create = mock.Mock()
         self._deleted = False
         self._transaction = transaction
+        self._last_use_time = last_use_time
 
     def __lt__(self, other):
         return id(self) < id(other)
 
+    @property
+    def last_use_time(self):
+        return self._last_use_time
+
     def exists(self):
         self._exists_checked = True
         return self._exists