Skip to content

Commit 747195b

Browse files
alex268alex268
committed
Merge branch 'develop' into auto_detect_query_type
2 parents 895b7f3 + 6a9f8a1 commit 747195b

File tree

6 files changed

+195
-53
lines changed

6 files changed

+195
-53
lines changed

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

Lines changed: 99 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,26 @@
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;
1013
import tech.ydb.jdbc.query.YdbQueryOptions;
14+
import tech.ydb.jdbc.settings.ParsedProperty;
1115
import tech.ydb.jdbc.settings.YdbClientProperties;
16+
import tech.ydb.jdbc.settings.YdbClientProperty;
1217
import tech.ydb.jdbc.settings.YdbConnectionProperties;
18+
import tech.ydb.jdbc.settings.YdbConnectionProperty;
1319
import tech.ydb.jdbc.settings.YdbOperationProperties;
1420
import tech.ydb.scheme.SchemeClient;
1521
import tech.ydb.table.TableClient;
22+
import tech.ydb.table.impl.PooledTableClient;
23+
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
1624

1725
/**
1826
*
@@ -22,20 +30,28 @@
2230
public class YdbContext implements AutoCloseable {
2331
private static final Logger LOGGER = Logger.getLogger(YdbContext.class.getName());
2432

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

2740
private final GrpcTransport grpcTransport;
28-
private final TableClient tableClient;
41+
private final PooledTableClient tableClient;
2942
private final SchemeClient schemeClient;
30-
3143
private final YdbQueryOptions queryOptions;
3244

33-
private YdbContext(YdbConfig config, GrpcTransport grpcTransport, TableClient tableClient) {
45+
private final boolean autoResizeSessionPool;
46+
private final AtomicInteger connectionsCount = new AtomicInteger();
47+
48+
private YdbContext(YdbConfig config, GrpcTransport transport, PooledTableClient tableClient, boolean autoResize) {
3449
this.config = config;
35-
this.grpcTransport = Objects.requireNonNull(grpcTransport);
50+
this.grpcTransport = Objects.requireNonNull(transport);
3651
this.tableClient = Objects.requireNonNull(tableClient);
37-
this.schemeClient = SchemeClient.newClient(grpcTransport).build();
52+
this.schemeClient = SchemeClient.newClient(transport).build();
3853
this.queryOptions = new YdbQueryOptions(config.getOperationProperties());
54+
this.autoResizeSessionPool = autoResize;
3955
}
4056

4157
public String getDatabase() {
@@ -58,6 +74,10 @@ public YdbQueryOptions getQueryOptions() {
5874
return queryOptions;
5975
}
6076

77+
public int getConnectionsCount() {
78+
return connectionsCount.get();
79+
}
80+
6181
public YdbOperationProperties getOperationProperties() {
6282
return config.getOperationProperties();
6383
}
@@ -73,20 +93,90 @@ public void close() {
7393
}
7494
}
7595

96+
public void register() {
97+
int actual = connectionsCount.incrementAndGet();
98+
int maxSize = tableClient.sessionPoolStats().getMaxSize();
99+
if (autoResizeSessionPool && actual > maxSize - SESSION_POOL_RESIZE_THRESHOLD) {
100+
int newSize = maxSize + SESSION_POOL_RESIZE_STEP;
101+
if (maxSize == tableClient.sessionPoolStats().getMaxSize()) {
102+
tableClient.updatePoolMaxSize(newSize);
103+
}
104+
}
105+
}
106+
107+
public void deregister() {
108+
int actual = connectionsCount.decrementAndGet();
109+
int maxSize = tableClient.sessionPoolStats().getMaxSize();
110+
if (autoResizeSessionPool && maxSize > SESSION_POOL_RESIZE_STEP) {
111+
if (actual < maxSize - SESSION_POOL_RESIZE_STEP - 2 * SESSION_POOL_RESIZE_THRESHOLD) {
112+
int newSize = maxSize - SESSION_POOL_RESIZE_STEP;
113+
if (maxSize == tableClient.sessionPoolStats().getMaxSize()) {
114+
tableClient.updatePoolMaxSize(newSize);
115+
}
116+
}
117+
}
118+
}
119+
76120
public static YdbContext createContext(YdbConfig config) throws SQLException {
77121
try {
78122
YdbConnectionProperties connProps = config.getConnectionProperties();
79123
YdbClientProperties clientProps = config.getClientProperties();
80124

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

83-
GrpcTransport grpcTransport = connProps.toGrpcTransport();
84-
85-
TableClient tableClient = clientProps.toTableClient(grpcTransport);
127+
GrpcTransport grpcTransport = buildGrpcTransport(connProps);
128+
PooledTableClient.Builder tableClient = PooledTableClient.newClient(
129+
GrpcTableRpc.useTransport(grpcTransport)
130+
);
131+
boolean autoResize = buildTableClient(tableClient, clientProps);
86132

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

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

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

7779
<T extends RequestSettings<?>> T withDefaultTimeout(T settings) {
@@ -184,6 +186,7 @@ public void close() throws SQLException {
184186
commit(); // like Oracle
185187
executor.clearWarnings();
186188
state = null;
189+
ctx.deregister();
187190
}
188191

189192
@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)