File tree

6 files changed

+83
-9
lines changed

6 files changed

+83
-9
lines changed
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,16 @@ private enum SessionMode {
3737
READ_WRITE
3838
}
3939

40+
@VisibleForTesting final String clientId;
4041
@VisibleForTesting final SessionPool pool;
4142

43+
@VisibleForTesting
4244
DatabaseClientImpl(SessionPool pool) {
45+
this("", pool);
46+
}
47+
48+
DatabaseClientImpl(String clientId, SessionPool pool) {
49+
this.clientId = clientId;
4350
this.pool = pool;
4451
}
4552

Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
class MetricRegistryConstants {
2424

2525
// The label keys are used to uniquely identify timeseries.
26+
private static final LabelKey CLIENT_ID =
27+
LabelKey.create("client_id", "User defined database client id");
2628
private static final LabelKey DATABASE = LabelKey.create("database", "Target database");
2729
private static final LabelKey INSTANCE_ID =
2830
LabelKey.create("instance_id", "Name of the instance");
@@ -33,10 +35,10 @@ class MetricRegistryConstants {
3335
private static final LabelValue UNSET_LABEL = LabelValue.create(null);
3436

3537
static final ImmutableList<LabelKey> SPANNER_LABEL_KEYS =
36-
ImmutableList.of(DATABASE, INSTANCE_ID, LIBRARY_VERSION);
38+
ImmutableList.of(CLIENT_ID, DATABASE, INSTANCE_ID, LIBRARY_VERSION);
3739

3840
static final ImmutableList<LabelValue> SPANNER_DEFAULT_LABEL_VALUES =
39-
ImmutableList.of(UNSET_LABEL, UNSET_LABEL, UNSET_LABEL);
41+
ImmutableList.of(UNSET_LABEL, UNSET_LABEL, UNSET_LABEL, UNSET_LABEL);
4042

4143
/** Unit to represent counts. */
4244
static final String COUNT = "1";
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,24 @@ class SpannerImpl extends BaseService<SpannerOptions> implements Spanner {
6262
static final String QUERY = "CloudSpannerOperation.ExecuteQuery";
6363
static final String READ = "CloudSpannerOperation.ExecuteRead";
6464

65+
private static final Object CLIENT_ID_LOCK = new Object();
66+
67+
@GuardedBy("CLIENT_ID_LOCK")
68+
private static final Map<DatabaseId, Long> CLIENT_IDS = new HashMap<>();
69+
70+
private static String nextDatabaseClientId(DatabaseId databaseId) {
71+
synchronized (CLIENT_ID_LOCK) {
72+
Long id = CLIENT_IDS.get(databaseId);
73+
if (id == null) {
74+
id = 1L;
75+
} else {
76+
id++;
77+
}
78+
CLIENT_IDS.put(databaseId, id);
79+
return String.format("client-%d", id);
80+
}
81+
}
82+
6583
private final SpannerRpc gapicRpc;
6684

6785
@GuardedBy("this")
@@ -153,24 +171,26 @@ public DatabaseClient getDatabaseClient(DatabaseId db) {
153171
if (dbClients.containsKey(db)) {
154172
return dbClients.get(db);
155173
} else {
174+
String clientId = nextDatabaseClientId(db);
156175
List<LabelValue> labelValues =
157176
ImmutableList.of(
177+
LabelValue.create(clientId),
158178
LabelValue.create(db.getDatabase()),
159179
LabelValue.create(db.getInstanceId().getName()),
160180
LabelValue.create(GaxProperties.getLibraryVersion(getOptions().getClass())));
161181
SessionPool pool =
162182
SessionPool.createPool(
163183
getOptions(), SpannerImpl.this.getSessionClient(db), labelValues);
164-
DatabaseClientImpl dbClient = createDatabaseClient(pool);
184+
DatabaseClientImpl dbClient = createDatabaseClient(clientId, pool);
165185
dbClients.put(db, dbClient);
166186
return dbClient;
167187
}
168188
}
169189
}
170190

171191
@VisibleForTesting
172-
DatabaseClientImpl createDatabaseClient(SessionPool pool) {
173-
return new DatabaseClientImpl(pool);
192+
DatabaseClientImpl createDatabaseClient(String clientId, SessionPool pool) {
193+
return new DatabaseClientImpl(clientId, pool);
174194
}
175195

176196
@Override
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@ private static class SpannerWithClosedSessionsImpl extends SpannerImpl {
4545
}
4646

4747
@Override
48-
DatabaseClientImpl createDatabaseClient(SessionPool pool) {
49-
return new DatabaseClientWithClosedSessionImpl(pool);
48+
DatabaseClientImpl createDatabaseClient(String clientId, SessionPool pool) {
49+
return new DatabaseClientWithClosedSessionImpl(clientId, pool);
5050
}
5151
}
5252

@@ -58,8 +58,8 @@ public static class DatabaseClientWithClosedSessionImpl extends DatabaseClientIm
5858
private boolean invalidateNextSession = false;
5959
private boolean allowReplacing = true;
6060

61-
DatabaseClientWithClosedSessionImpl(SessionPool pool) {
62-
super(pool);
61+
DatabaseClientWithClosedSessionImpl(String clientId, SessionPool pool) {
62+
super(clientId, pool);
6363
}
6464

6565
/** Invalidate the next session that is checked out from the pool. */
Original file line numberDiff line numberDiff line change
@@ -1579,6 +1579,7 @@ public void testSessionMetrics() throws Exception {
15791579
FakeMetricRegistry metricRegistry = new FakeMetricRegistry();
15801580
List<LabelValue> labelValues =
15811581
Arrays.asList(
1582+
LabelValue.create("client1"),
15821583
LabelValue.create("database1"),
15831584
LabelValue.create("instance1"),
15841585
LabelValue.create("1.0.0"));
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.util.Collections;
3333
import java.util.HashMap;
3434
import java.util.Map;
35+
import java.util.UUID;
3536
import org.junit.After;
3637
import org.junit.Before;
3738
import org.junit.Test;
@@ -185,6 +186,49 @@ public void testSpannerClosed() throws InterruptedException {
185186
spanner4.close();
186187
}
187188

189+
@Test
190+
public void testClientId() {
191+
// Create a unique database id to be sure it has not yet been used in the lifetime of this JVM.
192+
String dbName =
193+
String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString());
194+
DatabaseId db = DatabaseId.of(dbName);
195+
196+
Mockito.when(spannerOptions.getTransportOptions())
197+
.thenReturn(GrpcTransportOptions.newBuilder().build());
198+
Mockito.when(spannerOptions.getSessionPoolOptions())
199+
.thenReturn(SessionPoolOptions.newBuilder().setMinSessions(0).build());
200+
201+
DatabaseClientImpl databaseClient = (DatabaseClientImpl) impl.getDatabaseClient(db);
202+
assertThat(databaseClient.clientId).isEqualTo("client-1");
203+
204+
// Get same db client again.
205+
DatabaseClientImpl databaseClient1 = (DatabaseClientImpl) impl.getDatabaseClient(db);
206+
assertThat(databaseClient1.clientId).isEqualTo(databaseClient.clientId);
207+
208+
// Get a db client for a different database.
209+
String dbName2 =
210+
String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString());
211+
DatabaseId db2 = DatabaseId.of(dbName2);
212+
DatabaseClientImpl databaseClient2 = (DatabaseClientImpl) impl.getDatabaseClient(db2);
213+
assertThat(databaseClient2.clientId).isEqualTo("client-1");
214+
215+
// Create a new Spanner instance. This will generate new database clients with new ids.
216+
try (Spanner spanner =
217+
SpannerOptions.newBuilder()
218+
.setProjectId("p1")
219+
.setCredentials(NoCredentials.getInstance())
220+
.build()
221+
.getService()) {
222+
223+
// Get a database client for the same database as the first database. As this goes through a
224+
// different Spanner instance with potentially different options, it will get a different
225+
// client
226+
// id.
227+
DatabaseClientImpl databaseClient3 = (DatabaseClientImpl) spanner.getDatabaseClient(db);
228+
assertThat(databaseClient3.clientId).isEqualTo("client-2");
229+
}
230+
}
231+
188232
private SpannerOptions createSpannerOptions() {
189233
return SpannerOptions.newBuilder()
190234
.setProjectId("[PROJECT]")

0 commit comments

Comments
 (0)