fix: revert using GetSession as ping

googleapis · Apr 2, 2020 · 6401cb9 · 6401cb9
1 parent d2db09f
commit 6401cb9
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 30 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java
@@ -101,10 +101,6 @@ public String getName() {
     return options;
   }
 
-  com.google.spanner.v1.Session get() {
-    return spanner.getRpc().getSession(name);
-  }
-
   @Override
   public long executePartitionedUpdate(Statement stmt) {
     setActive(null);

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java
@@ -841,7 +841,12 @@ public void prepareReadWriteTransaction() {
 
     private void keepAlive() {
       markUsed();
-      delegate.get();
+      try (ResultSet resultSet =
+          delegate
+              .singleUse(TimestampBound.ofMaxStaleness(60, TimeUnit.SECONDS))
+              .executeQuery(Statement.newBuilder("SELECT 1").build())) {
+        resultSet.next();
+      }
     }
 
     void markUsed() {

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -109,7 +109,6 @@
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteBatchDmlResponse;
 import com.google.spanner.v1.ExecuteSqlRequest;
-import com.google.spanner.v1.GetSessionRequest;
 import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.PartitionQueryRequest;
 import com.google.spanner.v1.PartitionReadRequest;
@@ -345,13 +344,6 @@ private void acquireAdministrativeRequestsRateLimiter() {
     }
   }
 
-  @Override
-  public Session getSession(String name) {
-    GetSessionRequest request = GetSessionRequest.newBuilder().setName(name).build();
-    GrpcCallContext context = newCallContext(null, projectName);
-    return get(spannerStub.getSessionCallable().futureCall(request, context));
-  }
-
   @Override
   public Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable String pageToken)
       throws SpannerException {

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java
@@ -172,8 +172,6 @@ interface StreamingCall {
     void cancel(@Nullable String message);
   }
 
-  Session getSession(String name);
-
   // Instance admin APIs.
   Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable String pageToken)
       throws SpannerException;

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java
@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
@@ -99,20 +100,23 @@ public void run() {
   }
 
   private SessionImpl setupMockSession(final SessionImpl session) {
-    when(session.get())
+    ReadContext mockContext = mock(ReadContext.class);
+    final ResultSet mockResult = mock(ResultSet.class);
+    when(session.singleUse(any(TimestampBound.class))).thenReturn(mockContext);
+    when(mockContext.executeQuery(any(Statement.class)))
         .thenAnswer(
-            new Answer<com.google.spanner.v1.Session>() {
+            new Answer<ResultSet>() {
               @Override
-              public com.google.spanner.v1.Session answer(InvocationOnMock invocation)
-                  throws Throwable {
+              public ResultSet answer(InvocationOnMock invocation) throws Throwable {
                 Integer currentValue = pingedSessions.get(session.getName());
                 if (currentValue == null) {
                   currentValue = 0;
                 }
                 pingedSessions.put(session.getName(), ++currentValue);
-                return com.google.spanner.v1.Session.getDefaultInstance();
+                return mockResult;
               }
             });
+    when(mockResult.next()).thenReturn(true);
     return session;
   }
 

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java
@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -142,16 +143,20 @@ public Void answer(InvocationOnMock invocation) throws Throwable {
   }
 
   private void setupSession(final SessionImpl session) {
-    when(session.get())
+    ReadContext mockContext = mock(ReadContext.class);
+    final ResultSet mockResult = mock(ResultSet.class);
+    when(session.singleUse(any(TimestampBound.class))).thenReturn(mockContext);
+    when(mockContext.executeQuery(any(Statement.class)))
         .thenAnswer(
-            new Answer<com.google.spanner.v1.Session>() {
+            new Answer<ResultSet>() {
+
               @Override
-              public com.google.spanner.v1.Session answer(InvocationOnMock invocation)
-                  throws Throwable {
+              public ResultSet answer(InvocationOnMock invocation) throws Throwable {
                 resetTransaction(session);
-                return com.google.spanner.v1.Session.getDefaultInstance();
+                return mockResult;
               }
             });
+    when(mockResult.next()).thenReturn(true);
     doAnswer(
             new Answer<ApiFuture<Empty>>() {
 

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
@@ -892,8 +892,8 @@ public void run() {
             })
         .when(sessionClient)
         .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
-    for (SessionImpl sess : new SessionImpl[] {session1, session2, session3}) {
-      when(sess.get()).thenReturn(com.google.spanner.v1.Session.getDefaultInstance());
+    for (SessionImpl session : sessions) {
+      mockKeepAlive(session);
     }
     FakeClock clock = new FakeClock();
     clock.currentTimeMillis = System.currentTimeMillis();
@@ -963,16 +963,18 @@ public void run() {
     session1.close();
     session2.close();
     runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
-    verify(session, never()).get();
+    verify(session, never()).singleUse(any(TimestampBound.class));
     runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
-    verify(session, times(2)).get();
+    verify(session, times(2)).singleUse(any(TimestampBound.class));
     clock.currentTimeMillis +=
         clock.currentTimeMillis + (options.getKeepAliveIntervalMinutes() + 5) * 60 * 1000;
     session1 = pool.getReadSession();
     session1.writeAtLeastOnce(new ArrayList<Mutation>());
     session1.close();
     runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
-    verify(session, times(options.getMinSessions())).get();
+    // The session pool only keeps MinSessions + MaxIdleSessions alive.
+    verify(session, times(options.getMinSessions() + options.getMaxIdleSessions()))
+        .singleUse(any(TimestampBound.class));
     pool.closeAsync().get(5L, TimeUnit.SECONDS);
   }
 
@@ -1642,6 +1644,7 @@ public Void call() {
   private void mockKeepAlive(Session session) {
     ReadContext context = mock(ReadContext.class);
     ResultSet resultSet = mock(ResultSet.class);
+    when(resultSet.next()).thenReturn(true, false);
     when(session.singleUse(any(TimestampBound.class))).thenReturn(context);
     when(context.executeQuery(any(Statement.class))).thenReturn(resultSet);
   }