Skip to content

Commit 6a9f8a1

Browse files
alex268alex268
authored
Merge pull request #21 from ydb-platform/session_pool_fixes
Session pool fixes
2 parents 72354b6 + 40c3cb9 commit 6a9f8a1

File tree

6 files changed

+195
-52
lines changed

6 files changed

+195
-52
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java‎

Lines changed: 99 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,25 @@
11
package tech.ydb.jdbc.context;
22

33
import java.sql.SQLException;
4+
import java.util.Map;
45
import java.util.Objects;
6+
import java.util.concurrent.atomic.AtomicInteger;
57
import java.util.logging.Level;
68
import java.util.logging.Logger;
79

810
import tech.ydb.core.grpc.GrpcTransport;
11+
import tech.ydb.core.grpc.GrpcTransportBuilder;
912
import tech.ydb.jdbc.exception.YdbConfigurationException;
13+
import tech.ydb.jdbc.settings.ParsedProperty;
1014
import tech.ydb.jdbc.settings.YdbClientProperties;
15+
import tech.ydb.jdbc.settings.YdbClientProperty;
1116
import tech.ydb.jdbc.settings.YdbConnectionProperties;
17+
import tech.ydb.jdbc.settings.YdbConnectionProperty;
1218
import tech.ydb.jdbc.settings.YdbOperationProperties;
1319
import tech.ydb.scheme.SchemeClient;
1420
import tech.ydb.table.TableClient;
21+
import tech.ydb.table.impl.PooledTableClient;
22+
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
1523

1624
/**
1725
*
@@ -21,17 +29,26 @@
2129
public class YdbContext implements AutoCloseable {
2230
private static final Logger LOGGER = Logger.getLogger(YdbContext.class.getName());
2331

32+
private static final int SESSION_POOL_DEFAULT_MIN_SIZE = 0;
33+
private static final int SESSION_POOL_DEFAULT_MAX_SIZE = 50;
34+
private static final int SESSION_POOL_RESIZE_STEP = 50;
35+
private static final int SESSION_POOL_RESIZE_THRESHOLD = 10;
36+
2437
private final YdbConfig config;
2538

2639
private final GrpcTransport grpcTransport;
27-
private final TableClient tableClient;
40+
private final PooledTableClient tableClient;
2841
private final SchemeClient schemeClient;
42+
private final boolean autoResizeSessionPool;
43+
44+
private final AtomicInteger connectionsCount = new AtomicInteger();
2945

30-
private YdbContext(YdbConfig config, GrpcTransport grpcTransport, TableClient tableClient) {
46+
private YdbContext(YdbConfig config, GrpcTransport transport, PooledTableClient tableClient, boolean autoResize) {
3147
this.config = config;
32-
this.grpcTransport = Objects.requireNonNull(grpcTransport);
48+
this.grpcTransport = Objects.requireNonNull(transport);
3349
this.tableClient = Objects.requireNonNull(tableClient);
34-
this.schemeClient = SchemeClient.newClient(grpcTransport).build();
50+
this.schemeClient = SchemeClient.newClient(transport).build();
51+
this.autoResizeSessionPool = autoResize;
3552
}
3653

3754
public String getDatabase() {
@@ -50,6 +67,10 @@ public String getUrl() {
5067
return config.getUrl();
5168
}
5269

70+
public int getConnectionsCount() {
71+
return connectionsCount.get();
72+
}
73+
5374
public YdbOperationProperties getOperationProperties() {
5475
return config.getOperationProperties();
5576
}
@@ -65,20 +86,90 @@ public void close() {
6586
}
6687
}
6788

89+
public void register() {
90+
int actual = connectionsCount.incrementAndGet();
91+
int maxSize = tableClient.sessionPoolStats().getMaxSize();
92+
if (autoResizeSessionPool && actual > maxSize - SESSION_POOL_RESIZE_THRESHOLD) {
93+
int newSize = maxSize + SESSION_POOL_RESIZE_STEP;
94+
if (maxSize == tableClient.sessionPoolStats().getMaxSize()) {
95+
tableClient.updatePoolMaxSize(newSize);
96+
}
97+
}
98+
}
99+
100+
public void deregister() {
101+
int actual = connectionsCount.decrementAndGet();
102+
int maxSize = tableClient.sessionPoolStats().getMaxSize();
103+
if (autoResizeSessionPool && maxSize > SESSION_POOL_RESIZE_STEP) {
104+
if (actual < maxSize - SESSION_POOL_RESIZE_STEP - 2 * SESSION_POOL_RESIZE_THRESHOLD) {
105+
int newSize = maxSize - SESSION_POOL_RESIZE_STEP;
106+
if (maxSize == tableClient.sessionPoolStats().getMaxSize()) {
107+
tableClient.updatePoolMaxSize(newSize);
108+
}
109+
}
110+
}
111+
}
112+
68113
public static YdbContext createContext(YdbConfig config) throws SQLException {
69114
try {
70115
YdbConnectionProperties connProps = config.getConnectionProperties();
71116
YdbClientProperties clientProps = config.getClientProperties();
72117

73118
LOGGER.log(Level.INFO, "Creating new YDB connection to {0}", connProps.getConnectionString());
74119

75-
GrpcTransport grpcTransport = connProps.toGrpcTransport();
76-
77-
TableClient tableClient = clientProps.toTableClient(grpcTransport);
120+
GrpcTransport grpcTransport = buildGrpcTransport(connProps);
121+
PooledTableClient.Builder tableClient = PooledTableClient.newClient(
122+
GrpcTableRpc.useTransport(grpcTransport)
123+
);
124+
boolean autoResize = buildTableClient(tableClient, clientProps);
78125

79-
return new YdbContext(config, grpcTransport, tableClient);
126+
return new YdbContext(config, grpcTransport, tableClient.build(), autoResize);
80127
} catch (Exception ex) {
81128
throw new YdbConfigurationException("Cannot connect to YDB", ex);
82129
}
83130
}
131+
132+
public static GrpcTransport buildGrpcTransport(YdbConnectionProperties props) {
133+
GrpcTransportBuilder builder = GrpcTransport.forConnectionString(props.getConnectionString());
134+
for (Map.Entry<YdbConnectionProperty<?>, ParsedProperty> entry : props.getParams().entrySet()) {
135+
if (entry.getValue() != null) {
136+
entry.getKey().getSetter().accept(builder, entry.getValue().getParsedValue());
137+
}
138+
}
139+
140+
if (props.hasStaticCredentials()) {
141+
builder = builder.withAuthProvider(props.getStaticCredentials());
142+
}
143+
144+
return builder.build();
145+
}
146+
147+
private static boolean buildTableClient(TableClient.Builder builder, YdbClientProperties props) {
148+
for (Map.Entry<YdbClientProperty<?>, ParsedProperty> entry : props.getParams().entrySet()) {
149+
if (entry.getValue() != null) {
150+
entry.getKey().getSetter().accept(builder, entry.getValue().getParsedValue());
151+
}
152+
}
153+
154+
ParsedProperty minSizeConfig = props.getProperty(YdbClientProperty.SESSION_POOL_SIZE_MIN);
155+
ParsedProperty maxSizeConfig = props.getProperty(YdbClientProperty.SESSION_POOL_SIZE_MAX);
156+
157+
if (minSizeConfig == null && maxSizeConfig == null) {
158+
return true;
159+
}
160+
161+
int minSize = SESSION_POOL_DEFAULT_MIN_SIZE;
162+
int maxSize = SESSION_POOL_DEFAULT_MAX_SIZE;
163+
164+
if (minSizeConfig != null) {
165+
minSize = Math.max(0, minSizeConfig.getParsedValue());
166+
maxSize = Math.max(maxSize, minSize);
167+
}
168+
if (maxSizeConfig != null) {
169+
maxSize = Math.max(minSize + 1, maxSizeConfig.getParsedValue());
170+
}
171+
172+
builder.sessionPoolSize(minSize, maxSize);
173+
return false;
174+
}
84175
}

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ public YdbConnectionImpl(YdbContext context) throws SQLException {
7373
int txLevel = ctx.getOperationProperties().getTransactionLevel();
7474
boolean txAutoCommit = ctx.getOperationProperties().isAutoCommit();
7575
this.state = YdbTxState.create(txLevel, txAutoCommit);
76+
77+
this.ctx.register();
7678
}
7779

7880
<T extends RequestSettings<?>> T withDefaultTimeout(T settings) {
@@ -185,6 +187,7 @@ public void close() throws SQLException {
185187
commit(); // like Oracle
186188
executor.clearWarnings();
187189
state = null;
190+
ctx.deregister();
188191
}
189192

190193
@Override

‎jdbc/src/main/java/tech/ydb/jdbc/settings/YdbClientProperties.java‎

Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,6 @@
44

55
import javax.annotation.Nullable;
66

7-
import tech.ydb.core.grpc.GrpcTransport;
8-
import tech.ydb.table.TableClient;
9-
107
public class YdbClientProperties {
118
private final Map<YdbClientProperty<?>, ParsedProperty> params;
129

@@ -22,29 +19,4 @@ public ParsedProperty getProperty(YdbClientProperty<?> property) {
2219
public Map<YdbClientProperty<?>, ParsedProperty> getParams() {
2320
return params;
2421
}
25-
26-
public TableClient toTableClient(GrpcTransport grpc) {
27-
TableClient.Builder builder = TableClient.newClient(grpc);
28-
for (Map.Entry<YdbClientProperty<?>, ParsedProperty> entry : params.entrySet()) {
29-
if (entry.getValue() != null) {
30-
entry.getKey().getSetter().accept(builder, entry.getValue().getParsedValue());
31-
}
32-
}
33-
34-
int minSize = 0;
35-
int maxSize = 50;
36-
37-
ParsedProperty minSizeConfig = params.get(YdbClientProperty.SESSION_POOL_SIZE_MIN);
38-
ParsedProperty maxSizeConfig = params.get(YdbClientProperty.SESSION_POOL_SIZE_MAX);
39-
40-
if (minSizeConfig != null) {
41-
minSize = Math.max(0, minSizeConfig.getParsedValue());
42-
maxSize = Math.min(maxSize, minSize);
43-
}
44-
if (maxSizeConfig != null) {
45-
maxSize = Math.max(minSize, maxSizeConfig.getParsedValue());
46-
}
47-
48-
return builder.sessionPoolSize(minSize, maxSize).build();
49-
}
5022
}

‎jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConnectionProperties.java‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
import javax.annotation.Nullable;
77

88
import tech.ydb.core.auth.StaticCredentials;
9-
import tech.ydb.core.grpc.GrpcTransport;
10-
import tech.ydb.core.grpc.GrpcTransportBuilder;
9+
1110

1211
public class YdbConnectionProperties {
1312
private final String safeURL;
@@ -42,18 +41,11 @@ public Map<YdbConnectionProperty<?>, ParsedProperty> getParams() {
4241
return params;
4342
}
4443

45-
public GrpcTransport toGrpcTransport() {
46-
GrpcTransportBuilder builder = GrpcTransport.forConnectionString(connectionString);
47-
for (Map.Entry<YdbConnectionProperty<?>, ParsedProperty> entry : params.entrySet()) {
48-
if (entry.getValue() != null) {
49-
entry.getKey().getSetter().accept(builder, entry.getValue().getParsedValue());
50-
}
51-
}
52-
53-
if (username != null && !username.isEmpty()) {
54-
builder = builder.withAuthProvider(new StaticCredentials(username, password));
55-
}
44+
public boolean hasStaticCredentials() {
45+
return username != null && !username.isEmpty();
46+
}
5647

57-
return builder.build();
48+
public StaticCredentials getStaticCredentials() {
49+
return new StaticCredentials(username, password);
5850
}
5951
}

‎jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,16 @@
33
import java.sql.Connection;
44
import java.sql.DriverManager;
55
import java.sql.SQLException;
6+
import java.util.ArrayDeque;
7+
import java.util.Deque;
68
import java.util.Properties;
79

810
import org.junit.jupiter.api.Assertions;
911
import org.junit.jupiter.api.Test;
1012
import org.junit.jupiter.api.extension.RegisterExtension;
1113

12-
import tech.ydb.jdbc.impl.YdbConnectionImpl;
1314
import tech.ydb.jdbc.context.YdbContext;
15+
import tech.ydb.jdbc.impl.YdbConnectionImpl;
1416
import tech.ydb.jdbc.impl.helper.JdbcUrlHelper;
1517
import tech.ydb.test.junit5.YdbHelperExtension;
1618

@@ -95,4 +97,87 @@ public void testContextCache() throws SQLException {
9597
Assertions.assertNotSame(ctx, unwrapped.getCtx());
9698
}
9799
}
100+
101+
@Test
102+
public void testResizeSessionPool() throws SQLException {
103+
String url = jdbcURL.build();
104+
try (Connection conn = DriverManager.getConnection(url)) {
105+
YdbContext ctx = conn.unwrap(YdbConnection.class).getCtx();
106+
107+
Assertions.assertEquals(1, ctx.getConnectionsCount());
108+
Assertions.assertEquals(50, ctx.getTableClient().sessionPoolStats().getMaxSize());
109+
110+
Deque<Connection> connections = new ArrayDeque<>();
111+
for (int i = 0; i < 39; i++) {
112+
connections.offer(DriverManager.getConnection(url));
113+
}
114+
115+
Assertions.assertEquals(40, ctx.getConnectionsCount());
116+
Assertions.assertEquals(50, ctx.getTableClient().sessionPoolStats().getMaxSize());
117+
118+
connections.add(DriverManager.getConnection(url));
119+
120+
Assertions.assertEquals(41, ctx.getConnectionsCount());
121+
Assertions.assertEquals(100, ctx.getTableClient().sessionPoolStats().getMaxSize());
122+
123+
for (int i = 0; i < 11; i++) {
124+
connections.poll().close();
125+
}
126+
127+
Assertions.assertEquals(30, ctx.getConnectionsCount());
128+
Assertions.assertEquals(100, ctx.getTableClient().sessionPoolStats().getMaxSize());
129+
130+
connections.poll().close();
131+
132+
Assertions.assertEquals(29, ctx.getConnectionsCount());
133+
Assertions.assertEquals(50, ctx.getTableClient().sessionPoolStats().getMaxSize());
134+
135+
for (Connection c: connections) {
136+
c.close();
137+
}
138+
139+
Assertions.assertEquals(1, ctx.getConnectionsCount());
140+
Assertions.assertEquals(50, ctx.getTableClient().sessionPoolStats().getMaxSize());
141+
}
142+
}
143+
144+
@Test
145+
public void testFixedSessionPool() throws SQLException {
146+
assertFixedSessionPool("sessionPoolSizeMin", "0", 50);
147+
assertFixedSessionPool("sessionPoolSizeMin", "-1", 50);
148+
assertFixedSessionPool("sessionPoolSizeMax", "0", 1);
149+
assertFixedSessionPool("sessionPoolSizeMax", "5", 5);
150+
}
151+
152+
private void assertFixedSessionPool(String arg, String value, int poolSize) throws SQLException {
153+
String url = jdbcURL.withArg(arg, value).build();
154+
try (Connection conn = DriverManager.getConnection(url)) {
155+
YdbContext ctx = conn.unwrap(YdbConnection.class).getCtx();
156+
157+
Assertions.assertEquals(1, ctx.getConnectionsCount());
158+
Assertions.assertEquals(poolSize, ctx.getTableClient().sessionPoolStats().getMaxSize());
159+
160+
Deque<Connection> connections = new ArrayDeque<>();
161+
for (int i = 0; i < poolSize * 2; i++) {
162+
connections.offer(DriverManager.getConnection(url));
163+
}
164+
165+
Assertions.assertEquals(1 + poolSize * 2, ctx.getConnectionsCount());
166+
Assertions.assertEquals(poolSize, ctx.getTableClient().sessionPoolStats().getMaxSize());
167+
168+
for (int i = 0; i < poolSize; i++) {
169+
connections.poll().close();
170+
}
171+
172+
Assertions.assertEquals(1 + poolSize, ctx.getConnectionsCount());
173+
Assertions.assertEquals(poolSize, ctx.getTableClient().sessionPoolStats().getMaxSize());
174+
175+
for (Connection c: connections) {
176+
c.close();
177+
}
178+
179+
Assertions.assertEquals(1, ctx.getConnectionsCount());
180+
Assertions.assertEquals(poolSize, ctx.getTableClient().sessionPoolStats().getMaxSize());
181+
}
182+
}
98183
}

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
<slf4j.version>1.7.36</slf4j.version>
2121
<junit.version>5.9.3</junit.version>
2222

23-
<ydb.sdk.version>2.1.4</ydb.sdk.version>
23+
<ydb.sdk.version>2.1.5</ydb.sdk.version>
2424
</properties>
2525

2626
<licenses>

0 commit comments

Comments
 (0)