Skip to content

Commit ee6c2e9

Browse files
robertvoinescu-workrobertvoinescu-work
committed
feat(Spanner.V1): Add operation metrics for ReliableStreamReader.
1 parent 61cffee commit ee6c2e9

6 files changed

Lines changed: 270 additions & 24 deletions

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1.Tests/SpannerBuiltInMetricsInstrumentedOperationStreamReaderTests.cs‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
using Google.Cloud.Spanner.V1.Internal.Logging;
16+
using Google.Protobuf.WellKnownTypes;
17+
using Grpc.Core;
18+
using System;
19+
using System.Linq;
20+
using System.Threading;
21+
using System.Threading.Tasks;
22+
using Xunit;
23+
using static Google.Cloud.Spanner.V1.Tests.MetricsCapture;
24+
using static Google.Cloud.Spanner.V1.Tests.SpannerBuiltInMetricsInterceptorTests;
25+
using static Google.Cloud.Spanner.V1.Tests.SpannerBuiltInMetricsWrapperTests;
26+
27+
namespace Google.Cloud.Spanner.V1.Tests;
28+
29+
public class SpannerBuiltInMetricsInstrumentedOperationStreamReaderTests
30+
{
31+
private static readonly SpannerBuiltInMetrics.ClientIdentity s_readerClientIdentity = SpannerBuiltInMetrics.Labeler.GenerateIdentity();
32+
private static readonly ExecuteSqlRequest s_request = new() { SessionAsSessionName = s_sessionName };
33+
34+
public static TheoryData<int[]> StreamItemCases => new()
35+
{
36+
{ [] },
37+
{ [1] },
38+
{ [1, 2, 3] },
39+
};
40+
41+
public static TheoryData<Exception, StatusCode> InnerExceptionCases => new()
42+
{
43+
{ new RpcException(new Status(StatusCode.NotFound, "Test")), StatusCode.NotFound },
44+
{ new RpcException(new Status(StatusCode.DeadlineExceeded, "Test")), StatusCode.DeadlineExceeded },
45+
{ new RpcException(new Status(StatusCode.Cancelled, "Cancelled")), StatusCode.Cancelled },
46+
{ new OperationCanceledException("Operation was cancelled."), StatusCode.Unknown },
47+
{ new InvalidOperationException("Test"), StatusCode.Unknown },
48+
};
49+
50+
[Theory]
51+
[MemberData(nameof(StreamItemCases))]
52+
public async Task ConsumedStream_NoException_EmitsMetrics(int[] values)
53+
{
54+
var (reader, stopwatch) = CreateReader(values);
55+
56+
var measurements = await RunWithMeterListenerAsync(() => ConsumeReaderAsync(reader));
57+
58+
AssertRecordedData(measurements, s_readerClientIdentity.Id, "ExecuteStreamingSql", StatusCode.OK.ToString());
59+
Assert.True(stopwatch.Stopped);
60+
}
61+
62+
[Theory]
63+
[MemberData(nameof(InnerExceptionCases))]
64+
public async Task ConsumedStream_InnerException_EmitsMetrics(Exception exception, StatusCode expectedStatus)
65+
{
66+
var (reader, stopwatch) = CreateReader(exception: exception);
67+
68+
var measurements = await RunWithMeterListenerAsync(async () =>
69+
{
70+
var thrown = await Assert.ThrowsAnyAsync<Exception>(() => ConsumeReaderAsync(reader));
71+
Assert.Same(exception, thrown);
72+
});
73+
74+
AssertRecordedData(measurements, s_readerClientIdentity.Id, "ExecuteStreamingSql", expectedStatus.ToString());
75+
Assert.True(stopwatch.Stopped);
76+
}
77+
78+
[Fact]
79+
public async Task ClosedBeforeCompletion_Emits()
80+
{
81+
var (reader, stopwatch) = CreateReader();
82+
83+
var measurements = await RunWithMeterListenerAsync(async () =>
84+
{
85+
await reader.NextAsync(CancellationToken.None);
86+
reader.Close();
87+
});
88+
89+
AssertRecordedData(measurements, s_readerClientIdentity.Id, "ExecuteStreamingSql", StatusCode.Unknown.ToString());
90+
Assert.True(stopwatch.Stopped);
91+
}
92+
93+
[Fact]
94+
public async Task DisposedBeforeCompletion_Emits()
95+
{
96+
var (reader, stopwatch) = CreateReader();
97+
98+
var measurements = await RunWithMeterListenerAsync(async () =>
99+
{
100+
await reader.NextAsync(CancellationToken.None);
101+
reader.Dispose();
102+
});
103+
104+
AssertRecordedData(measurements, s_readerClientIdentity.Id, "ExecuteStreamingSql", StatusCode.Unknown.ToString());
105+
Assert.True(stopwatch.Stopped);
106+
}
107+
108+
[Fact]
109+
public async Task ClosedAfterCompletion_EmitsOnce()
110+
{
111+
var (reader, stopwatch) = CreateReader();
112+
113+
var measurements = await RunWithMeterListenerAsync(async () =>
114+
{
115+
// Only the first call should result in metrics being emitted
116+
await ConsumeReaderAsync(reader);
117+
reader.Close();
118+
reader.Dispose();
119+
});
120+
121+
AssertRecordedData(measurements, s_readerClientIdentity.Id, "ExecuteStreamingSql", StatusCode.OK.ToString());
122+
Assert.True(stopwatch.Stopped);
123+
}
124+
125+
[Fact]
126+
public void Disposed_DisposesInnerStream()
127+
{
128+
var inner = new DisposableAsyncStreamReader();
129+
var instrumentedStream = new SpannerBuiltInMetrics.InstrumentedOperationStreamReader(
130+
inner,
131+
s_stopwatchProvider.StartNew(),
132+
"ExecuteStreamingSql",
133+
s_request,
134+
s_readerClientIdentity);
135+
136+
instrumentedStream.Dispose();
137+
138+
Assert.True(inner.Disposed);
139+
}
140+
141+
/// <summary>
142+
/// Creates a reader over a stream with one <see cref="PartialResultSet"/> per value in <paramref name="values"/>.
143+
/// </summary>
144+
private static (ReliableStreamReader reader, FakeStopwatch stopwatch) CreateReader(int[] values = null, Exception exception = null)
145+
{
146+
var stopwatch = (FakeStopwatch) s_stopwatchProvider.StartNew();
147+
var resultSets = (values ?? [1]).Select(value => new PartialResultSet
148+
{
149+
// Spanner encodes INT64 values as strings.
150+
Values = { Value.ForString(value.ToString()) },
151+
Metadata = new ResultSetMetadata
152+
{
153+
RowType = new StructType
154+
{
155+
Fields = { new StructType.Types.Field { Name = "col1", Type = new Type { Code = TypeCode.Int64 } } }
156+
}
157+
}
158+
}).ToList();
159+
160+
var fakeStream = new FakeAsyncStreamReader<PartialResultSet>(resultSets, exception);
161+
var instrumentedStream = new SpannerBuiltInMetrics.InstrumentedOperationStreamReader(
162+
fakeStream,
163+
stopwatch,
164+
"ExecuteStreamingSql",
165+
s_request,
166+
s_readerClientIdentity);
167+
var reader = new ReliableStreamReader(instrumentedStream, Logger.DefaultLogger);
168+
return (reader, stopwatch);
169+
}
170+
171+
private static async Task ConsumeReaderAsync(ReliableStreamReader reader)
172+
{
173+
while (await reader.NextAsync(CancellationToken.None) != null)
174+
{
175+
}
176+
}
177+
178+
private sealed class DisposableAsyncStreamReader : IAsyncStreamReader<PartialResultSet>, IDisposable
179+
{
180+
public bool Disposed { get; private set; }
181+
182+
public PartialResultSet Current => null;
183+
184+
public Task<bool> MoveNext(CancellationToken cancellationToken) => Task.FromResult(false);
185+
186+
public void Dispose() => Disposed = true;
187+
}
188+
}

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1.Tests/SpannerBuiltInMetricsWrapperTests.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,7 @@ public async Task RecordsMetricsOnException(
454454
private static List<Measurement> FilterMeasurements(IEnumerable<Measurement> measurements, string clientUid) =>
455455
measurements.Where(m => m.GetTag("client_uid") == clientUid).ToList();
456456

457-
private static void AssertRecordedData(IEnumerable<Measurement> measurements, string clientUid, string expectedMethod, string expectedStatus)
457+
internal static void AssertRecordedData(IEnumerable<Measurement> measurements, string clientUid, string expectedMethod, string expectedStatus)
458458
{
459459
var scoped = FilterMeasurements(measurements, clientUid);
460460
Assert.Single(scoped, m => m.Name == "operation_count");

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/ManagedTransaction.cs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -720,7 +720,20 @@ internal ReliableStreamReader ExecuteReadOrQueryStreamReader(ReadOrQueryRequest
720720
_client.MaybeApplyRouteToLeaderHeader(ref callSettings, _transactionOptions.ModeCase);
721721
MaybeApplyDirectedReadOptions(request.UnderlyingRequest);
722722

723-
ResultStream stream = new ResultStream(_client, request, this, callSettings);
723+
Grpc.Core.IAsyncStreamReader<PartialResultSet> stream = new ResultStream(_client, request, this, callSettings);
724+
// Only clients built by SpannerClientBuilder are wrapped in a BuiltInMetricsWrapper and record operation metrics.
725+
if (_client is SpannerBuiltInMetrics.BuiltInMetricsWrapper metricsWrapper)
726+
{
727+
string methodName = request.IsQuery
728+
? nameof(SpannerClient.ExecuteStreamingSql)
729+
: nameof(SpannerClient.StreamingRead);
730+
stream = new SpannerBuiltInMetrics.InstrumentedOperationStreamReader(
731+
stream,
732+
SpannerBuiltInMetrics.DefaultStopwatchProvider.Instance.StartNew(),
733+
methodName,
734+
request,
735+
metricsWrapper._clientIdentity);
736+
}
724737
return new ReliableStreamReader(stream, _client.Settings.Logger);
725738
}
726739

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerBuiltInMetrics.BuiltInMetricsWrapper.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,8 @@ private TResponse ExecuteWithInstrumentation<TResponse>(string method, object re
187187
private void RecordMetrics(string method, object request, StatusCode status, double latencyMs)
188188
{
189189
var provider = request as IDatabaseNameProvider;
190-
RecordOperationMetrics(latencyMs, method, provider, status, _clientIdentity);
190+
var labels = Labeler.GetLabels(method, provider, status, _clientIdentity);
191+
RecordOperationMetrics(latencyMs, labels);
191192
}
192193
}
193194
}

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerBuiltInMetrics.InstrumentedOperationStreamReader.cs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
using Grpc.Core;
16+
using System;
17+
using System.Collections.Generic;
18+
using System.Threading.Tasks;
19+
20+
namespace Google.Cloud.Spanner.V1;
21+
22+
internal static partial class SpannerBuiltInMetrics
23+
{
24+
/// <summary>
25+
/// Decorates the retrying stream of <see cref="PartialResultSet"/> that backs a <see cref="ReliableStreamReader"/>
26+
/// to record operation metrics.
27+
/// </summary>
28+
/// <remarks>
29+
/// The inner stream retries transient errors itself, so each gRPC stream it opens is a separate attempt,
30+
/// recorded by <see cref="InstrumentedAttemptStreamReader{TResponse}"/>. This decorator records the operation as a whole.
31+
/// The operation completes when the inner stream is fully consumed, fails with an error that won't be retried,
32+
/// or is disposed before completion, which <see cref="ReliableStreamReader.Close"/> does.
33+
/// </remarks>
34+
internal sealed class InstrumentedOperationStreamReader : InstrumentedStreamReaderBase<PartialResultSet>, IDisposable
35+
{
36+
internal InstrumentedOperationStreamReader(
37+
IAsyncStreamReader<PartialResultSet> inner,
38+
IStopwatch stopwatch,
39+
string methodName,
40+
IDatabaseNameProvider dbNameProvider,
41+
ClientIdentity clientIdentity)
42+
: base(inner, stopwatch, methodName, dbNameProvider, clientIdentity)
43+
{
44+
}
45+
46+
/// <inheritdoc/>
47+
public void Dispose()
48+
{
49+
// The stream was disposed before it was fully consumed. If metrics were already emitted, this is a no-op.
50+
// Recording operation metrics is synchronous, so the returned task is already complete.
51+
_ = EmitMetricsAsync(StatusCode.Unknown);
52+
if (Inner is IDisposable disposable)
53+
{
54+
disposable.Dispose();
55+
}
56+
}
57+
58+
/// <inheritdoc/>
59+
private protected override Task RecordAsync(double latencyMs, KeyValuePair<string, object>[] labels)
60+
{
61+
RecordOperationMetrics(latencyMs, labels);
62+
return Task.CompletedTask;
63+
}
64+
}
65+
}

‎apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerBuiltInMetrics.cs‎

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -86,27 +86,6 @@ internal static void RecordAttemptMetrics(double latencyMs, KeyValuePair<string,
8686
}
8787
}
8888

89-
/// <summary>
90-
/// Records operation count and latency metrics.
91-
/// </summary>
92-
/// <param name="latencyMs">The elapsed duration of the operation in milliseconds.</param>
93-
/// <param name="methodName">The name of the gRPC method invoked.</param>
94-
/// <param name="dbNameProvider">The provider holding database context details.</param>
95-
/// <param name="status">The resolved status of the operation.</param>
96-
/// <param name="clientIdentity">The identity context for the client executing the operation.</param>
97-
internal static void RecordOperationMetrics(double latencyMs, string methodName, IDatabaseNameProvider dbNameProvider, StatusCode status, ClientIdentity clientIdentity)
98-
{
99-
try
100-
{
101-
var labels = Labeler.GetLabels(methodName, dbNameProvider, status, clientIdentity);
102-
RecordOperationMetrics(latencyMs, labels);
103-
}
104-
catch
105-
{
106-
// Silently swallow exceptions.
107-
}
108-
}
109-
11089
/// <summary>
11190
/// Records operation count and latency metrics.
11291
/// </summary>

0 commit comments

Comments
 (0)