[shardingsphere] branch master updated: Remove PostgreSQLConnectionContext when channel inactive (#11087)

[zh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zhangliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git


The following commit(s) were added to refs/heads/master by this push:
     new fda105e  Remove PostgreSQLConnectionContext when channel inactive (#11087)
fda105e is described below

commit fda105e345efaaa4830e570765a7c25e63978ab6
Author: 吴伟杰 <wu...@apache.org>
AuthorDate: Thu Jul 1 14:40:16 2021 +0800

    Remove PostgreSQLConnectionContext when channel inactive (#11087)
---
 .../frontend/postgresql/PostgreSQLFrontendEngine.java   |  2 ++
 .../postgresql/PostgreSQLFrontendEngineTest.java        | 17 +++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngine.java b/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngine.java
index 46a9f81..0e15258 100644
--- a/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngine.java
+++ b/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngine.java
@@ -28,6 +28,7 @@ import org.apache.shardingsphere.proxy.frontend.command.CommandExecuteEngine;
 import org.apache.shardingsphere.proxy.frontend.context.FrontendContext;
 import org.apache.shardingsphere.proxy.frontend.postgresql.authentication.PostgreSQLAuthenticationEngine;
 import org.apache.shardingsphere.proxy.frontend.postgresql.command.PostgreSQLCommandExecuteEngine;
+import org.apache.shardingsphere.proxy.frontend.postgresql.command.PostgreSQLConnectionContextRegistry;
 import org.apache.shardingsphere.proxy.frontend.spi.DatabaseProtocolFrontendEngine;
 
 /**
@@ -47,6 +48,7 @@ public final class PostgreSQLFrontendEngine implements DatabaseProtocolFrontendE
     @Override
     public void release(final BackendConnection backendConnection) {
         PostgreSQLBinaryStatementRegistry.getInstance().unregister(backendConnection.getConnectionId());
+        PostgreSQLConnectionContextRegistry.getInstance().remove(backendConnection.getConnectionId());
     }
     
     @Override
diff --git a/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngineTest.java b/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngineTest.java
index d1f3da0..ef995bc 100644
--- a/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngineTest.java
+++ b/shardingsphere-proxy/shardingsphere-proxy-frontend/shardingsphere-proxy-frontend-postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/PostgreSQLFrontendEngineTest.java
@@ -17,16 +17,23 @@
 
 package org.apache.shardingsphere.proxy.frontend.postgresql;
 
+import lombok.SneakyThrows;
 import org.apache.shardingsphere.db.protocol.postgresql.packet.command.query.binary.PostgreSQLBinaryStatementRegistry;
 import org.apache.shardingsphere.infra.database.type.dialect.PostgreSQLDatabaseType;
 import org.apache.shardingsphere.proxy.backend.communication.jdbc.connection.BackendConnection;
 import org.apache.shardingsphere.proxy.frontend.executor.ConnectionThreadExecutorGroup;
+import org.apache.shardingsphere.proxy.frontend.postgresql.command.PostgreSQLConnectionContext;
+import org.apache.shardingsphere.proxy.frontend.postgresql.command.PostgreSQLConnectionContextRegistry;
 import org.junit.Test;
 
+import java.lang.reflect.Field;
+import java.util.concurrent.ConcurrentMap;
+
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -46,11 +53,21 @@ public final class PostgreSQLFrontendEngineTest {
         when(backendConnection.getConnectionId()).thenReturn(connectionId);
         PostgreSQLBinaryStatementRegistry registry = PostgreSQLBinaryStatementRegistry.getInstance();
         registry.register(connectionId);
+        PostgreSQLConnectionContextRegistry.getInstance().get(connectionId);
         assertNotNull(registry.get(connectionId));
         PostgreSQLFrontendEngine frontendEngine = new PostgreSQLFrontendEngine();
         ConnectionThreadExecutorGroup.getInstance().register(connectionId);
         ConnectionThreadExecutorGroup.getInstance().unregisterAndAwaitTermination(connectionId);
         frontendEngine.release(backendConnection);
         assertNull(registry.get(connectionId));
+        assertTrue(getConnectionContexts().isEmpty());
+    }
+    
+    @SuppressWarnings("unchecked")
+    @SneakyThrows
+    private ConcurrentMap<Integer, PostgreSQLConnectionContext> getConnectionContexts() {
+        Field field = PostgreSQLConnectionContextRegistry.class.getDeclaredField("connectionContexts");
+        field.setAccessible(true);
+        return (ConcurrentMap<Integer, PostgreSQLConnectionContext>) field.get(PostgreSQLConnectionContextRegistry.getInstance());
     }
 }