fix(sql-execute): add an user config and add concurrent control for session creating (#2020)

Author: yhilmare

server/odc-core/src/main/java/com/oceanbase/odc/core/session/BaseConnectionSessionManager.java

@@ -56,6 +56,11 @@ public ConnectionSession start(ConnectionSessionFactory factory) {
         try {
             doStoreSession(session);
         } catch (Throwable e) {
+            try {
+                session.expire();
+            } catch (Exception e1) {
+                log.warn("Failed to expire session, session storage failed, sessId={}", session.getId(), e1);
+            }
             log.warn("Failed to store a session, session={}", session, e);
             try {
                 onCreateFailed(session, e);

server/odc-service/src/main/java/com/oceanbase/odc/service/dispatch/RequestDispatcher.java

@@ -77,13 +77,11 @@ public DispatchResponse forward(@NonNull String ip, @NonNull Integer port) throw
     }
 
     public DispatchResponse forward(@NonNull String ip, @NonNull Integer port, HttpServletRequest request,
-            ByteArrayOutputStream requestBody)
-            throws IOException {
+            ByteArrayOutputStream requestBody) throws IOException {
         String hostUrl = getHostUrl(ip, port);
         return forward(hostUrl, request, requestBody);
     }
 
-
     private DispatchResponse forward(String hostUrl, HttpServletRequest request, ByteArrayOutputStream requestBody)
             throws IOException {
         Verify.notNull(request, "HttpServletRequest");
@@ -96,7 +94,6 @@ private DispatchResponse forward(String hostUrl, HttpServletRequest request, Byt
         return forward(hostUrl, method, requestUrl, headers, requestBody.toByteArray());
     }
 
-
     public String getHostUrl(@NonNull String ip, @NonNull Integer port) {
         return String.format("%s://%s:%s", PROTOCAL, ip, port);
     }
@@ -118,8 +115,6 @@ public String getRequestUrlByRequest(HttpServletRequest request) {
         return uriBuilder.toString();
     }
 
-
-
     public DispatchResponse forward(@NonNull String hostUrl, @NonNull HttpMethod method,
             @NonNull String requestUri, @NonNull HttpHeaders headers, byte[] requestBody) throws IOException {
         verifyAndReduceTtl(headers);

server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java

@@ -447,7 +447,6 @@ private List<SqlTuple> generateSqlTuple(List<OffsetString> sqls, ConnectionSessi
             }
             if (Objects.isNull(request.getQueryLimit())
                     || !request.ifAddROWID()
-                    || !sessionProperties.isAddInternalRowId()
                     || session.getDialectType() != DialectType.OB_ORACLE) {
                 return target;
             }

server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectSessionService.java

@@ -25,7 +25,10 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
@@ -146,9 +149,9 @@ public class ConnectSessionService {
     private CloudMetadataClient cloudMetadataClient;
     @Autowired
     private UserConfigFacade userConfigFacade;
-
     @Autowired
     private StateHostGenerator stateHostGenerator;
+    private final Map<String, Lock> sessionId2Lock = new ConcurrentHashMap<>();
 
     @PostConstruct
     public void init() {
@@ -158,6 +161,7 @@ public void init() {
         this.connectionSessionManager = new DefaultConnectionSessionManager(
                 new DefaultTaskManager("connection-session-management"), repository);
         this.connectionSessionManager.addListener(new SessionLimitListener(limitService));
+        this.connectionSessionManager.addListener(new SessionLockRemoveListener(this.sessionId2Lock));
         this.connectionSessionManager.enableAsyncRefreshSessionManager();
         this.connectionSessionManager.addSessionValidator(
                 new SessionValidatorPredicate(sessionProperties.getTimeoutMins(), TimeUnit.MINUTES));
@@ -333,9 +337,25 @@ public ConnectionSession nullSafeGet(@NotNull String sessionId, boolean autoCrea
             if (!autoCreate || !StringUtils.equals(req.getFrom(), stateHostGenerator.getHost())) {
                 throw new NotFoundException(ResourceType.ODC_SESSION, "ID", sessionId);
             }
-            session = create(req);
-            ConnectionSessionUtil.setConsoleSessionResetFlag(session, true);
-            return session;
+            Lock lock = this.sessionId2Lock.computeIfAbsent(sessionId, s -> new ReentrantLock());
+            try {
+                if (!lock.tryLock(10, TimeUnit.SECONDS)) {
+                    throw new IllegalStateException("Session is creating, please wait and retry later");
+                }
+            } catch (Exception e) {
+                throw new IllegalStateException(e);
+            }
+            try {
+                session = connectionSessionManager.getSession(sessionId);
+                if (session != null) {
+                    return session;
+                }
+                session = create(req);
+                ConnectionSessionUtil.setConsoleSessionResetFlag(session, true);
+                return session;
+            } finally {
+                lock.unlock();
+            }
         }
         if (!Objects.equals(ConnectionSessionUtil.getUserId(session), authenticationFacade.currentUserId())) {
             throw new NotFoundException(ResourceType.ODC_SESSION, "ID", sessionId);

server/odc-service/src/main/java/com/oceanbase/odc/service/session/SessionLockRemoveListener.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) 2023 OceanBase.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.oceanbase.odc.service.session;
+
+import java.util.Map;
+import java.util.concurrent.locks.Lock;
+
+import com.oceanbase.odc.core.session.ConnectionSession;
+import com.oceanbase.odc.core.session.DefaultSessionEventListener;
+
+/**
+ * {@link SessionLockRemoveListener}
+ *
+ * @author yh263208
+ * @date 2024-03-22 16:59
+ * @since ODC_release_4.2.4
+ * @see DefaultSessionEventListener
+ */
+public class SessionLockRemoveListener extends DefaultSessionEventListener {
+
+    private final Map<String, Lock> sessionId2Lock;
+
+    public SessionLockRemoveListener(Map<String, Lock> sessionId2Lock) {
+        this.sessionId2Lock = sessionId2Lock;
+    }
+
+    @Override
+    public void onExpire(ConnectionSession session) {
+        this.sessionId2Lock.remove(session.getId());
+    }
+
+}

server/odc-service/src/main/java/com/oceanbase/odc/service/session/SessionProperties.java

@@ -116,12 +116,6 @@ public class SessionProperties {
     @Value("${odc.session.timeout-mins:480}")
     private long timeoutMins = 480;
 
-    /**
-     * 用于控制用户在 SQL 执行时是否改写增加 rowid 列
-     */
-    @Value("${odc.session.sql-execute.add-internal-rowid:true}")
-    private boolean addInternalRowId = true;
-
     /**
      * Timeout for querying full link trace
      */

server/odc-service/src/main/resources/config-meta/user-config-meta.yml

@@ -161,3 +161,13 @@
   default_value: 'false'
   nullable: false
   description: ''
+- id: 17
+  category: 'user'
+  key: 'odc.sqlexecute.default.addInternalRowId'
+  type: 'boolean'
+  max_value: ~
+  min_value: ~
+  allowed_values: [ 'true','false' ]
+  default_value: 'true'
+  nullable: false
+  description: 'Whether the query is being rewritten in ob-oracle mode and increases the RowId.'