Skip to content

Commit 3e2b45f

Browse files
agrawal-siddharthagrawal-siddharth
committed
feat: populate client_stats in AppendRowsRequest
1 parent 6bb477b commit 3e2b45f

2 files changed

Lines changed: 163 additions & 0 deletions

File tree

‎java-bigquerystorage/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1156,6 +1156,8 @@ private void appendLoop() {
11561156
boolean firstRequestForTableOrSchemaSwitch = true;
11571157

11581158
while (!waitingQueueDrained()) {
1159+
long queuedRequestsCount = 0;
1160+
ClientStats.WindowStats windowStats = null;
11591161
this.lock.lock();
11601162
try {
11611163
hasMessageInWaitingQueue.await(100, TimeUnit.MILLISECONDS);
@@ -1203,6 +1205,9 @@ private void appendLoop() {
12031205
localQueue.addLast(requestWrapper);
12041206
healthCheckMetrics.updateRequestsSent(requestWrapper.messageSize);
12051207
}
1208+
queuedRequestsCount =
1209+
inflightRequestQueue.size(); // waitingRequestQueue is empty so can be omitted here
1210+
windowStats = createWindowStats();
12061211
} catch (InterruptedException e) {
12071212
log.warning(
12081213
"Interrupted while waiting for message. Stream: "
@@ -1289,6 +1294,12 @@ private void appendLoop() {
12891294
}
12901295
firstRequestForTableOrSchemaSwitch = false;
12911296

1297+
ClientStats clientStats = createClientStats(wrapper, queuedRequestsCount, windowStats);
1298+
windowStats = null;
1299+
if (!clientStats.equals(ClientStats.getDefaultInstance())) {
1300+
originalRequestBuilder.setClientStats(clientStats);
1301+
}
1302+
12921303
requestProfilerHook.startOperation(
12931304
RequestProfiler.OperationName.RESPONSE_LATENCY, requestUniqueId);
12941305

@@ -1307,6 +1318,65 @@ private void appendLoop() {
13071318
cleanupConnectionAndRequests(/* avoidBlocking= */ false);
13081319
}
13091320

1321+
private ClientStats createClientStats(
1322+
AppendRequestAndResponse requestWrapper,
1323+
long queuedRequestsCount,
1324+
@Nullable ClientStats.WindowStats windowStats) {
1325+
ClientStats.Builder clientStatsBuilder = ClientStats.newBuilder();
1326+
ClientStats.RequestStats requestStats = createRequestStats(requestWrapper, queuedRequestsCount);
1327+
if (!requestStats.equals(ClientStats.RequestStats.getDefaultInstance())) {
1328+
clientStatsBuilder.setRequestStats(requestStats);
1329+
}
1330+
if (windowStats != null && !windowStats.equals(ClientStats.WindowStats.getDefaultInstance())) {
1331+
clientStatsBuilder.setWindowStats(windowStats);
1332+
}
1333+
return clientStatsBuilder.build();
1334+
}
1335+
1336+
private ClientStats.RequestStats createRequestStats(
1337+
AppendRequestAndResponse requestWrapper, long queuedRequestsCount) {
1338+
ClientStats.RequestStats.Builder requestStatsBuilder = ClientStats.RequestStats.newBuilder();
1339+
if (requestWrapper.requestSendTimeStamp != null) {
1340+
requestStatsBuilder.setSendTimeMillis(requestWrapper.requestSendTimeStamp.toEpochMilli());
1341+
}
1342+
if (queuedRequestsCount > 0) {
1343+
requestStatsBuilder.setQueuedRequestsCount(queuedRequestsCount);
1344+
}
1345+
return requestStatsBuilder.build();
1346+
}
1347+
1348+
@GuardedBy("lock")
1349+
private ClientStats.WindowStats createWindowStats() {
1350+
ClientStats.WindowStats.Builder windowStatsBuilder = ClientStats.WindowStats.newBuilder();
1351+
if (healthCheckMetrics.windowedMilliLatencyMax != 0) {
1352+
windowStatsBuilder.setMaxResponseLatencyMillis(healthCheckMetrics.windowedMilliLatencyMax);
1353+
}
1354+
long avgResponseLatencyMillis =
1355+
healthCheckMetrics.windowedResponsesAcked > 0
1356+
? healthCheckMetrics.windowedMilliLatencySum / healthCheckMetrics.windowedResponsesAcked
1357+
: 0;
1358+
if (avgResponseLatencyMillis != 0) {
1359+
windowStatsBuilder.setAvgResponseLatencyMillis(avgResponseLatencyMillis);
1360+
}
1361+
if (healthCheckMetrics.windowedMilliResponseWaitTimeMax != 0) {
1362+
windowStatsBuilder.setLongestWaitNoResponseMillis(
1363+
healthCheckMetrics.windowedMilliResponseWaitTimeMax);
1364+
}
1365+
if (healthCheckMetrics.windowedRequestsSent != 0) {
1366+
windowStatsBuilder.setRequestsSentCount(healthCheckMetrics.windowedRequestsSent);
1367+
}
1368+
if (healthCheckMetrics.windowedResponsesAcked != 0) {
1369+
windowStatsBuilder.setResponsesReceivedCount(healthCheckMetrics.windowedResponsesAcked);
1370+
}
1371+
if (healthCheckMetrics.windowedRequestsSentBytes != 0) {
1372+
windowStatsBuilder.setBytesSentCount(healthCheckMetrics.windowedRequestsSentBytes);
1373+
}
1374+
long windowStartTimeEpochMillis = healthCheckMetrics.healthCheckTimeStamp.toEpochMilli();
1375+
windowStatsBuilder.setWindowStartTimeEpochMillis(windowStartTimeEpochMillis);
1376+
windowStatsBuilder.setWindowMillis(Instant.now().toEpochMilli() - windowStartTimeEpochMillis);
1377+
return windowStatsBuilder.build();
1378+
}
1379+
13101380
@Nullable
13111381
private AppendRowsSchema getSchema(AppendRowsRequest request) {
13121382
if (request.hasProtoRows()) {

‎java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1612,6 +1612,17 @@ void testInflightRetryCountHealthMetricExactlyOnce() throws Exception {
16121612
assertTrue(healthCheckFields.responseCodes.containsKey(Status.Code.OK.value()));
16131613
assertEquals(3, healthCheckFields.responseCodes.get(Status.Code.OK.value()));
16141614
assertEquals("projects/p1/datasets/d1/tables/t1/streams/s1", healthCheckFields.streamName);
1615+
1616+
AppendRowsRequest secondRequest = testBigQueryWrite.getAppendRequests().get(1);
1617+
assertThat(secondRequest.getClientStats().getRequestStats().hasQueuedRequestsCount()).isTrue();
1618+
assertThat(secondRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1619+
.isEqualTo(2);
1620+
// When both requests are retried together in localQueue, only the first request in localQueue
1621+
// should have window_stats populated.
1622+
AppendRowsRequest retryFirstInLocalQueue = testBigQueryWrite.getAppendRequests().get(2);
1623+
AppendRowsRequest retrySecondInLocalQueue = testBigQueryWrite.getAppendRequests().get(3);
1624+
assertThat(retryFirstInLocalQueue.getClientStats().hasWindowStats()).isTrue();
1625+
assertThat(retrySecondInLocalQueue.getClientStats().hasWindowStats()).isFalse();
16151626
}
16161627

16171628
@Test
@@ -1722,6 +1733,88 @@ void testEarliestSendTime_retryScenario() throws Exception {
17221733

17231734
// Once exhausted, state should be null
17241735
assertThat(connectionWorker.getEarliestSendTime()).isNull();
1736+
1737+
AppendRowsRequest initialRequest = testBigQueryWrite.getAppendRequests().get(0);
1738+
AppendRowsRequest retryRequest = testBigQueryWrite.getAppendRequests().get(1);
1739+
assertThat(initialRequest.getClientStats().getRequestStats().hasSendTimeMillis()).isTrue();
1740+
assertThat(retryRequest.getClientStats().getRequestStats().hasSendTimeMillis()).isTrue();
1741+
assertThat(retryRequest.getClientStats().getRequestStats().getSendTimeMillis())
1742+
.isAtLeast(initialRequest.getClientStats().getRequestStats().getSendTimeMillis());
1743+
assertThat(initialRequest.getClientStats().getRequestStats().hasQueuedRequestsCount())
1744+
.isTrue();
1745+
assertThat(initialRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1746+
.isEqualTo(1);
1747+
assertThat(retryRequest.getClientStats().getRequestStats().hasQueuedRequestsCount()).isTrue();
1748+
assertThat(retryRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1749+
.isEqualTo(1);
1750+
}
1751+
}
1752+
1753+
@Test
1754+
void testWindowStats() throws Exception {
1755+
long beforeCreationMillis = Instant.now().toEpochMilli();
1756+
ProtoSchema schema1 = createProtoSchema("foo");
1757+
StreamWriter sw1 =
1758+
StreamWriter.newBuilder(TEST_STREAM_1, client)
1759+
.setLocation("us")
1760+
.setWriterSchema(schema1)
1761+
.build();
1762+
try (ConnectionWorker connectionWorker = createConnectionWorker()) {
1763+
long afterCreationMillis = Instant.now().toEpochMilli();
1764+
int msecResponseDelay = 300;
1765+
testBigQueryWrite.setResponseSleep(java.time.Duration.ofMillis(msecResponseDelay));
1766+
testBigQueryWrite.addResponse(createAppendResponse(0));
1767+
testBigQueryWrite.addResponse(createAppendResponse(1));
1768+
1769+
// First request: no prior responses or wait times have been recorded yet.
1770+
ApiFuture<AppendRowsResponse> future1 =
1771+
sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"0"}), 0);
1772+
future1.get();
1773+
1774+
AppendRowsRequest firstRequest = testBigQueryWrite.getAppendRequests().get(0);
1775+
long sizePerRequest = firstRequest.getProtoRows().getSerializedSize();
1776+
assertThat(firstRequest.getClientStats().hasWindowStats()).isTrue();
1777+
ClientStats.WindowStats firstWindowStats = firstRequest.getClientStats().getWindowStats();
1778+
assertThat(firstWindowStats.hasMaxResponseLatencyMillis()).isFalse();
1779+
assertThat(firstWindowStats.hasAvgResponseLatencyMillis()).isFalse();
1780+
assertThat(firstWindowStats.hasLongestWaitNoResponseMillis()).isFalse();
1781+
assertThat(firstWindowStats.hasRequestsSentCount()).isTrue();
1782+
assertThat(firstWindowStats.getRequestsSentCount()).isEqualTo(1);
1783+
assertThat(firstWindowStats.hasResponsesReceivedCount()).isFalse();
1784+
assertThat(firstWindowStats.hasBytesSentCount()).isTrue();
1785+
assertThat(firstWindowStats.getBytesSentCount()).isEqualTo(sizePerRequest);
1786+
assertThat(firstWindowStats.hasWindowStartTimeEpochMillis()).isTrue();
1787+
assertThat(firstWindowStats.getWindowStartTimeEpochMillis()).isAtLeast(beforeCreationMillis);
1788+
assertThat(firstWindowStats.getWindowStartTimeEpochMillis()).isAtMost(afterCreationMillis);
1789+
assertThat(firstWindowStats.hasWindowMillis()).isTrue();
1790+
assertThat(firstWindowStats.getWindowMillis()).isAtLeast(0L);
1791+
1792+
// Second request: sent after firstRequest waited in-flight (>100ms) and received a response
1793+
// (>=300ms latency), so all expected WindowStats fields are now non-zero and populated.
1794+
ApiFuture<AppendRowsResponse> future2 =
1795+
sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"1"}), 1);
1796+
future2.get();
1797+
1798+
AppendRowsRequest secondRequest = testBigQueryWrite.getAppendRequests().get(1);
1799+
assertThat(secondRequest.getClientStats().hasWindowStats()).isTrue();
1800+
ClientStats.WindowStats secondWindowStats = secondRequest.getClientStats().getWindowStats();
1801+
assertThat(secondWindowStats.hasMaxResponseLatencyMillis()).isTrue();
1802+
assertThat(secondWindowStats.getMaxResponseLatencyMillis()).isAtLeast(msecResponseDelay);
1803+
assertThat(secondWindowStats.hasAvgResponseLatencyMillis()).isTrue();
1804+
assertThat(secondWindowStats.getAvgResponseLatencyMillis()).isAtLeast(msecResponseDelay);
1805+
assertThat(secondWindowStats.hasLongestWaitNoResponseMillis()).isTrue();
1806+
assertThat(secondWindowStats.getLongestWaitNoResponseMillis()).isGreaterThan(0L);
1807+
assertThat(secondWindowStats.hasRequestsSentCount()).isTrue();
1808+
assertThat(secondWindowStats.getRequestsSentCount()).isEqualTo(2);
1809+
assertThat(secondWindowStats.hasResponsesReceivedCount()).isTrue();
1810+
assertThat(secondWindowStats.getResponsesReceivedCount()).isEqualTo(1);
1811+
assertThat(secondWindowStats.hasBytesSentCount()).isTrue();
1812+
assertThat(secondWindowStats.getBytesSentCount()).isEqualTo(2 * sizePerRequest);
1813+
assertThat(secondWindowStats.hasWindowStartTimeEpochMillis()).isTrue();
1814+
assertThat(secondWindowStats.getWindowStartTimeEpochMillis())
1815+
.isEqualTo(firstWindowStats.getWindowStartTimeEpochMillis());
1816+
assertThat(secondWindowStats.hasWindowMillis()).isTrue();
1817+
assertThat(secondWindowStats.getWindowMillis()).isAtLeast(msecResponseDelay);
17251818
}
17261819
}
17271820
}

0 commit comments

Comments
 (0)