Skip to content

Commit 51dee39

Browse files
feat(Spanner.V1): Add wrapper for built in metrics.
1 parent 934b57d commit 51dee39

3 files changed

Lines changed: 622 additions & 3 deletions

File tree

Lines changed: 374 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,374 @@
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 NSubstitute;
16+
using System;
17+
using System.Diagnostics.Metrics;
18+
using System.Linq;
19+
using System.Reflection;
20+
using System.Threading.Tasks;
21+
using Xunit;
22+
23+
namespace Google.Cloud.Spanner.V1.Tests;
24+
25+
public class SpannerBuiltInMetricsWrapperTests
26+
{
27+
private static readonly Grpc.Core.Status s_timeoutStatus = new(Grpc.Core.StatusCode.DeadlineExceeded, "timeout");
28+
private static readonly Exception s_rpcException = new Grpc.Core.RpcException(s_timeoutStatus);
29+
30+
[Fact]
31+
public void WrapperOverridesAllCoreVirtualMethods()
32+
{
33+
var baseType = typeof(SpannerClient);
34+
var wrapperType = typeof(SpannerBuiltInMetrics).GetNestedType("BuiltInMetricsWrapper", BindingFlags.NonPublic);
35+
var methods = baseType.GetMethods(BindingFlags.Public | BindingFlags.Instance | BindingFlags.DeclaredOnly);
36+
37+
var missedMethods = new System.Collections.Generic.List<string>();
38+
foreach (var m in methods)
39+
{
40+
if (m.IsVirtual && !m.IsFinal)
41+
{
42+
var parameters = m.GetParameters();
43+
if (parameters.Length == 2 &&
44+
typeof(Google.Protobuf.IMessage).IsAssignableFrom(parameters[0].ParameterType) &&
45+
parameters[1].ParameterType == typeof(Google.Api.Gax.Grpc.CallSettings))
46+
{
47+
var paramTypes = parameters.Select(p => p.ParameterType).ToArray();
48+
var wm = wrapperType.GetMethod(m.Name, paramTypes);
49+
if (wm == null || wm.DeclaringType != wrapperType)
50+
{
51+
missedMethods.Add(m.Name);
52+
}
53+
}
54+
}
55+
}
56+
Assert.Empty(missedMethods);
57+
}
58+
59+
60+
[Fact]
61+
public async Task ConsecutiveCalls_RecordsMultipleMetrics()
62+
{
63+
var mockClient = Substitute.For<SpannerClient>();
64+
mockClient.ExecuteSqlAsync(Arg.Any<ExecuteSqlRequest>(), null).Returns(Task.FromResult(new ResultSet()));
65+
66+
var wrapper = new SpannerBuiltInMetrics.BuiltInMetricsWrapper(mockClient);
67+
68+
using var meterListener = CreateMeterListener(out var recordedData);
69+
meterListener.Start();
70+
71+
var request = new ExecuteSqlRequest();
72+
73+
// Execute exactly 3 repeated payload invocations.
74+
await wrapper.ExecuteSqlAsync(request, null);
75+
await wrapper.ExecuteSqlAsync(request, null);
76+
await wrapper.ExecuteSqlAsync(request, null);
77+
78+
meterListener.RecordObservableInstruments();
79+
80+
// Evaluate numerical accrual constraints robustly.
81+
Assert.Equal(3, recordedData.CountRecordings);
82+
Assert.Equal(3, recordedData.LatencyRecordings);
83+
Assert.Equal(nameof(SpannerClient.ExecuteSqlAsync), recordedData.Method);
84+
Assert.Equal(Grpc.Core.StatusCode.OK.ToString(), recordedData.Status);
85+
}
86+
87+
88+
[Theory]
89+
[MemberData(nameof(SuccessfulResponses))]
90+
public async Task RecordsTelemetryOnSuccess(
91+
Action<SpannerClient> setupClient,
92+
Func<SpannerClient, Task> invoke,
93+
string expectedMethod,
94+
string expectedStatus)
95+
{
96+
var mockClient = Substitute.For<SpannerClient>();
97+
setupClient(mockClient);
98+
99+
var wrapper = new SpannerBuiltInMetrics.BuiltInMetricsWrapper(mockClient);
100+
101+
using var meterListener = CreateMeterListener(out var recordedData);
102+
meterListener.Start();
103+
104+
await invoke(wrapper);
105+
106+
meterListener.RecordObservableInstruments();
107+
AssertRecordedData(recordedData, expectedMethod, expectedStatus);
108+
}
109+
110+
[Theory]
111+
[MemberData(nameof(UnsuccessfulResponses))]
112+
public async Task RecordsTelemetryOnException(
113+
Action<SpannerClient> setupClient,
114+
Func<SpannerClient, Task> invoke,
115+
string expectedMethod,
116+
string expectedStatus,
117+
Exception expectedException)
118+
{
119+
var mockClient = Substitute.For<SpannerClient>();
120+
setupClient(mockClient);
121+
122+
var wrapper = new SpannerBuiltInMetrics.BuiltInMetricsWrapper(mockClient);
123+
124+
using var meterListener = CreateMeterListener(out var recordedData);
125+
meterListener.Start();
126+
127+
var thrownException = await Assert.ThrowsAnyAsync<Exception>(() => invoke(wrapper));
128+
Assert.Equal(expectedException.GetType(), thrownException.GetType());
129+
130+
meterListener.RecordObservableInstruments();
131+
AssertRecordedData(recordedData, expectedMethod, expectedStatus);
132+
}
133+
134+
public static TheoryData<Action<SpannerClient>, Func<SpannerClient, Task>, string, string> SuccessfulResponses => new()
135+
{
136+
{
137+
c => c.CreateSession(Arg.Any<CreateSessionRequest>(), null).Returns(new Session()),
138+
w => { w.CreateSession(new CreateSessionRequest(), null); return Task.CompletedTask; },
139+
nameof(SpannerClient.CreateSession),
140+
Grpc.Core.StatusCode.OK.ToString()
141+
},
142+
{
143+
c => c.CreateSessionAsync(Arg.Any<CreateSessionRequest>(), null).Returns(Task.FromResult(new Session())),
144+
w => w.CreateSessionAsync(new CreateSessionRequest(), null),
145+
nameof(SpannerClient.CreateSessionAsync),
146+
Grpc.Core.StatusCode.OK.ToString()
147+
},
148+
149+
{
150+
c => c.BatchCreateSessions(Arg.Any<BatchCreateSessionsRequest>(), null).Returns(new BatchCreateSessionsResponse()),
151+
w => { w.BatchCreateSessions(new BatchCreateSessionsRequest(), null); return Task.CompletedTask; },
152+
nameof(SpannerClient.BatchCreateSessions),
153+
Grpc.Core.StatusCode.OK.ToString()
154+
},
155+
{
156+
c => c.BatchCreateSessionsAsync(Arg.Any<BatchCreateSessionsRequest>(), null).Returns(Task.FromResult(new BatchCreateSessionsResponse())),
157+
w => w.BatchCreateSessionsAsync(new BatchCreateSessionsRequest(), null),
158+
nameof(SpannerClient.BatchCreateSessionsAsync),
159+
Grpc.Core.StatusCode.OK.ToString()
160+
},
161+
162+
{
163+
c => c.GetSession(Arg.Any<GetSessionRequest>(), null).Returns(new Session()),
164+
w => { w.GetSession(new GetSessionRequest(), null); return Task.CompletedTask; },
165+
nameof(SpannerClient.GetSession),
166+
Grpc.Core.StatusCode.OK.ToString()
167+
},
168+
{
169+
c => c.GetSessionAsync(Arg.Any<GetSessionRequest>(), null).Returns(Task.FromResult(new Session())),
170+
w => w.GetSessionAsync(new GetSessionRequest(), null),
171+
nameof(SpannerClient.GetSessionAsync),
172+
Grpc.Core.StatusCode.OK.ToString()
173+
},
174+
175+
{
176+
c => c.ListSessions(Arg.Any<ListSessionsRequest>(), null).Returns((Google.Api.Gax.PagedEnumerable<ListSessionsResponse, Session>)null),
177+
w => { w.ListSessions(new ListSessionsRequest(), null); return Task.CompletedTask; },
178+
nameof(SpannerClient.ListSessions),
179+
Grpc.Core.StatusCode.OK.ToString()
180+
},
181+
{
182+
c => c.ListSessionsAsync(Arg.Any<ListSessionsRequest>(), null).Returns((Google.Api.Gax.PagedAsyncEnumerable<ListSessionsResponse, Session>)null),
183+
w => { w.ListSessionsAsync(new ListSessionsRequest(), null); return Task.CompletedTask; },
184+
nameof(SpannerClient.ListSessionsAsync),
185+
Grpc.Core.StatusCode.OK.ToString()
186+
},
187+
188+
{
189+
c => c.DeleteSession(Arg.Any<DeleteSessionRequest>(), null),
190+
w => { w.DeleteSession(new DeleteSessionRequest(), null); return Task.CompletedTask; },
191+
nameof(SpannerClient.DeleteSession),
192+
Grpc.Core.StatusCode.OK.ToString()
193+
},
194+
{
195+
c => c.DeleteSessionAsync(Arg.Any<DeleteSessionRequest>(), null).Returns(Task.CompletedTask),
196+
w => w.DeleteSessionAsync(new DeleteSessionRequest(), null),
197+
nameof(SpannerClient.DeleteSessionAsync),
198+
Grpc.Core.StatusCode.OK.ToString()
199+
},
200+
201+
{
202+
c => c.ExecuteSql(Arg.Any<ExecuteSqlRequest>(), null).Returns(new ResultSet()),
203+
w => { w.ExecuteSql(new ExecuteSqlRequest(), null); return Task.CompletedTask; },
204+
nameof(SpannerClient.ExecuteSql),
205+
Grpc.Core.StatusCode.OK.ToString()
206+
},
207+
{
208+
c => c.ExecuteSqlAsync(Arg.Any<ExecuteSqlRequest>(), null).Returns(Task.FromResult(new ResultSet())),
209+
w => w.ExecuteSqlAsync(new ExecuteSqlRequest(), null),
210+
nameof(SpannerClient.ExecuteSqlAsync),
211+
Grpc.Core.StatusCode.OK.ToString()
212+
},
213+
214+
{
215+
c => c.ExecuteBatchDml(Arg.Any<ExecuteBatchDmlRequest>(), null).Returns(new ExecuteBatchDmlResponse()),
216+
w => { w.ExecuteBatchDml(new ExecuteBatchDmlRequest(), null); return Task.CompletedTask; },
217+
nameof(SpannerClient.ExecuteBatchDml),
218+
Grpc.Core.StatusCode.OK.ToString()
219+
},
220+
{
221+
c => c.ExecuteBatchDmlAsync(Arg.Any<ExecuteBatchDmlRequest>(), null).Returns(Task.FromResult(new ExecuteBatchDmlResponse())),
222+
w => w.ExecuteBatchDmlAsync(new ExecuteBatchDmlRequest(), null),
223+
nameof(SpannerClient.ExecuteBatchDmlAsync),
224+
Grpc.Core.StatusCode.OK.ToString()
225+
},
226+
227+
{
228+
c => c.Read(Arg.Any<ReadRequest>(), null).Returns(new ResultSet()),
229+
w => { w.Read(new ReadRequest(), null); return Task.CompletedTask; },
230+
nameof(SpannerClient.Read),
231+
Grpc.Core.StatusCode.OK.ToString()
232+
},
233+
{
234+
c => c.ReadAsync(Arg.Any<ReadRequest>(), null).Returns(Task.FromResult(new ResultSet())),
235+
w => w.ReadAsync(new ReadRequest(), null),
236+
nameof(SpannerClient.ReadAsync),
237+
Grpc.Core.StatusCode.OK.ToString()
238+
},
239+
240+
{
241+
c => c.BeginTransaction(Arg.Any<BeginTransactionRequest>(), null).Returns(new Transaction()),
242+
w => { w.BeginTransaction(new BeginTransactionRequest(), null); return Task.CompletedTask; },
243+
nameof(SpannerClient.BeginTransaction),
244+
Grpc.Core.StatusCode.OK.ToString()
245+
},
246+
{
247+
c => c.BeginTransactionAsync(Arg.Any<BeginTransactionRequest>(), null).Returns(Task.FromResult(new Transaction())),
248+
w => w.BeginTransactionAsync(new BeginTransactionRequest(), null),
249+
nameof(SpannerClient.BeginTransactionAsync),
250+
Grpc.Core.StatusCode.OK.ToString()
251+
},
252+
253+
{
254+
c => c.Commit(Arg.Any<CommitRequest>(), null).Returns(new CommitResponse()),
255+
w => { w.Commit(new CommitRequest(), null); return Task.CompletedTask; },
256+
nameof(SpannerClient.Commit),
257+
Grpc.Core.StatusCode.OK.ToString()
258+
},
259+
{
260+
c => c.CommitAsync(Arg.Any<CommitRequest>(), null).Returns(Task.FromResult(new CommitResponse())),
261+
w => w.CommitAsync(new CommitRequest(), null),
262+
nameof(SpannerClient.CommitAsync),
263+
Grpc.Core.StatusCode.OK.ToString()
264+
},
265+
266+
{
267+
c => c.Rollback(Arg.Any<RollbackRequest>(), null),
268+
w => { w.Rollback(new RollbackRequest(), null); return Task.CompletedTask; },
269+
nameof(SpannerClient.Rollback),
270+
Grpc.Core.StatusCode.OK.ToString()
271+
},
272+
{
273+
c => c.RollbackAsync(Arg.Any<RollbackRequest>(), null).Returns(Task.CompletedTask),
274+
w => w.RollbackAsync(new RollbackRequest(), null),
275+
nameof(SpannerClient.RollbackAsync),
276+
Grpc.Core.StatusCode.OK.ToString()
277+
},
278+
279+
{
280+
c => c.PartitionQuery(Arg.Any<PartitionQueryRequest>(), null).Returns(new PartitionResponse()),
281+
w => { w.PartitionQuery(new PartitionQueryRequest(), null); return Task.CompletedTask; },
282+
nameof(SpannerClient.PartitionQuery),
283+
Grpc.Core.StatusCode.OK.ToString()
284+
},
285+
{
286+
c => c.PartitionQueryAsync(Arg.Any<PartitionQueryRequest>(), null).Returns(Task.FromResult(new PartitionResponse())),
287+
w => w.PartitionQueryAsync(new PartitionQueryRequest(), null),
288+
nameof(SpannerClient.PartitionQueryAsync),
289+
Grpc.Core.StatusCode.OK.ToString()
290+
},
291+
292+
{
293+
c => c.PartitionRead(Arg.Any<PartitionReadRequest>(), null).Returns(new PartitionResponse()),
294+
w => { w.PartitionRead(new PartitionReadRequest(), null); return Task.CompletedTask; },
295+
nameof(SpannerClient.PartitionRead),
296+
Grpc.Core.StatusCode.OK.ToString()
297+
},
298+
{
299+
c => c.PartitionReadAsync(Arg.Any<PartitionReadRequest>(), null).Returns(Task.FromResult(new PartitionResponse())),
300+
w => w.PartitionReadAsync(new PartitionReadRequest(), null),
301+
nameof(SpannerClient.PartitionReadAsync),
302+
Grpc.Core.StatusCode.OK.ToString()
303+
},
304+
};
305+
306+
public static TheoryData<Action<SpannerClient>, Func<SpannerClient, Task>, string, string, Exception> UnsuccessfulResponses => new()
307+
{
308+
{
309+
client => client.ExecuteSql(Arg.Any<ExecuteSqlRequest>(), null).Returns(x => throw s_rpcException),
310+
wrapper => { wrapper.ExecuteSql(new ExecuteSqlRequest { Session = "projects/p/instances/i/databases/d/sessions/s" }, null); return Task.CompletedTask; },
311+
nameof(SpannerClient.ExecuteSql),
312+
Grpc.Core.StatusCode.DeadlineExceeded.ToString(),
313+
s_rpcException
314+
},
315+
{
316+
client => client.CommitAsync(Arg.Any<CommitRequest>(), null).Returns(Task.FromException<CommitResponse>(new InvalidOperationException("boom"))),
317+
wrapper => wrapper.CommitAsync(new CommitRequest { Session = "projects/p/instances/i/databases/d/sessions/s" }, null),
318+
nameof(SpannerClient.CommitAsync),
319+
Grpc.Core.StatusCode.Unknown.ToString().ToUpperInvariant(),
320+
new InvalidOperationException("boom")
321+
}
322+
};
323+
324+
private class RecordedData
325+
{
326+
public int CountRecordings { get; set; }
327+
public int LatencyRecordings { get; set; }
328+
public string Method { get; set; }
329+
public string Status { get; set; }
330+
}
331+
332+
private static MeterListener CreateMeterListener(out RecordedData recordedData)
333+
{
334+
var data = new RecordedData();
335+
recordedData = data;
336+
337+
var listener = new MeterListener();
338+
listener.InstrumentPublished = (instrument, l) =>
339+
{
340+
if (instrument.Name == "operation_count" || instrument.Name == "operation_latencies")
341+
{
342+
l.EnableMeasurementEvents(instrument);
343+
}
344+
};
345+
346+
listener.SetMeasurementEventCallback<long>((instrument, measurement, tags, state) =>
347+
{
348+
if (instrument.Name == "operation_count")
349+
{
350+
data.CountRecordings++;
351+
data.Status = tags.ToArray().FirstOrDefault(t => t.Key == "status").Value?.ToString();
352+
data.Method = tags.ToArray().FirstOrDefault(t => t.Key == "method").Value?.ToString();
353+
}
354+
});
355+
356+
listener.SetMeasurementEventCallback<double>((instrument, measurement, tags, state) =>
357+
{
358+
if (instrument.Name == "operation_latencies")
359+
{
360+
data.LatencyRecordings++;
361+
}
362+
});
363+
364+
return listener;
365+
}
366+
367+
private static void AssertRecordedData(RecordedData data, string expectedMethod, string expectedStatus)
368+
{
369+
Assert.Equal(1, data.CountRecordings);
370+
Assert.Equal(1, data.LatencyRecordings);
371+
Assert.Equal(expectedMethod, data.Method);
372+
Assert.Equal(expectedStatus, data.Status);
373+
}
374+
}

0 commit comments

Comments
 (0)