Skip to content

Commit decfbb0

Browse files
sebright2sebright2
committed
Refactor the code for overriding the Census implementation in tests.
This commit adds a method to set the CensusStatsModule on AbstractServerImplBuilder and AbstractManagedChannelImplBuilder. It is simpler to set the whole CensusStatsModule than pass in all necessary OpenCensus implementation objects.
1 parent 997fdd1 commit decfbb0

File tree

7 files changed

+110
-100
lines changed

7 files changed

+110
-100
lines changed

‎core/src/main/java/io/grpc/internal/AbstractManagedChannelImplBuilder.java‎

Lines changed: 8 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,6 @@
3232
import io.grpc.NameResolver;
3333
import io.grpc.NameResolverProvider;
3434
import io.grpc.PickFirstBalancerFactory;
35-
import io.opencensus.stats.Stats;
36-
import io.opencensus.stats.StatsRecorder;
37-
import io.opencensus.tags.Tagger;
38-
import io.opencensus.tags.Tags;
39-
import io.opencensus.tags.propagation.TagContextBinarySerializer;
4035
import io.opencensus.trace.Tracing;
4136
import java.net.SocketAddress;
4237
import java.net.URI;
@@ -155,13 +150,7 @@ protected final int maxInboundMessageSize() {
155150
private boolean tracingEnabled = true;
156151

157152
@Nullable
158-
private Tagger tagger;
159-
160-
@Nullable
161-
private TagContextBinarySerializer tagCtxSerializer;
162-
163-
@Nullable
164-
private StatsRecorder statsRecorder;
153+
private CensusStatsModule censusStatsOverride;
165154

166155
protected AbstractManagedChannelImplBuilder(String target) {
167156
this.target = Preconditions.checkNotNull(target, "target");
@@ -296,11 +285,8 @@ public final T idleTimeout(long value, TimeUnit unit) {
296285
* Override the default stats implementation.
297286
*/
298287
@VisibleForTesting
299-
protected final T statsImplementation(
300-
Tagger tagger, TagContextBinarySerializer tagCtxSerializer, StatsRecorder statsRecorder) {
301-
this.tagger = tagger;
302-
this.tagCtxSerializer = tagCtxSerializer;
303-
this.statsRecorder = statsRecorder;
288+
protected final T overrideCensusStatsModule(CensusStatsModule censusStats) {
289+
this.censusStatsOverride = censusStats;
304290
return thisT();
305291
}
306292

@@ -358,24 +344,13 @@ final List<ClientInterceptor> getEffectiveInterceptors() {
358344
List<ClientInterceptor> effectiveInterceptors =
359345
new ArrayList<ClientInterceptor>(this.interceptors);
360346
if (statsEnabled) {
361-
Tagger tagger = this.tagger != null ? this.tagger : Tags.getTagger();
362-
TagContextBinarySerializer tagCtxSerializer =
363-
this.tagCtxSerializer != null
364-
? this.tagCtxSerializer
365-
: Tags.getTagPropagationComponent().getBinarySerializer();
366-
StatsRecorder statsRecorder =
367-
this.statsRecorder != null ? this.statsRecorder : Stats.getStatsRecorder();
368-
CensusStatsModule censusStats =
369-
new CensusStatsModule(
370-
tagger,
371-
tagCtxSerializer,
372-
statsRecorder,
373-
GrpcUtil.STOPWATCH_SUPPLIER,
374-
true,
375-
recordStats);
347+
CensusStatsModule censusStats = this.censusStatsOverride;
348+
if (censusStats == null) {
349+
censusStats = new CensusStatsModule(GrpcUtil.STOPWATCH_SUPPLIER, true);
350+
}
376351
// First interceptor runs last (see ClientInterceptors.intercept()), so that no
377352
// other interceptor can override the tracer factory we set in CallOptions.
378-
effectiveInterceptors.add(0, censusStats.getClientInterceptor());
353+
effectiveInterceptors.add(0, censusStats.getClientInterceptor(recordStats));
379354
}
380355
if (tracingEnabled) {
381356
CensusTracingModule censusTracing =

‎core/src/main/java/io/grpc/internal/AbstractServerImplBuilder.java‎

Lines changed: 8 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,6 @@
3434
import io.grpc.ServerServiceDefinition;
3535
import io.grpc.ServerStreamTracer;
3636
import io.grpc.ServerTransportFilter;
37-
import io.opencensus.stats.Stats;
38-
import io.opencensus.stats.StatsRecorder;
39-
import io.opencensus.tags.Tagger;
40-
import io.opencensus.tags.Tags;
41-
import io.opencensus.tags.propagation.TagContextBinarySerializer;
4237
import io.opencensus.trace.Tracing;
4338
import java.util.ArrayList;
4439
import java.util.Collections;
@@ -100,13 +95,7 @@ public List<ServerServiceDefinition> getServices() {
10095
CompressorRegistry compressorRegistry = DEFAULT_COMPRESSOR_REGISTRY;
10196

10297
@Nullable
103-
private Tagger tagger;
104-
105-
@Nullable
106-
private TagContextBinarySerializer tagCtxSerializer;
107-
108-
@Nullable
109-
private StatsRecorder statsRecorder;
98+
private CensusStatsModule censusStatsOverride;
11099

111100
private boolean statsEnabled = true;
112101
private boolean recordStats = true;
@@ -193,13 +182,8 @@ public final T compressorRegistry(CompressorRegistry registry) {
193182
* Override the default stats implementation.
194183
*/
195184
@VisibleForTesting
196-
protected T statsImplementation(
197-
final Tagger tagger,
198-
TagContextBinarySerializer tagCtxSerializer,
199-
StatsRecorder statsRecorder) {
200-
this.tagger = tagger;
201-
this.tagCtxSerializer = tagCtxSerializer;
202-
this.statsRecorder = statsRecorder;
185+
protected T overrideCensusStatsModule(CensusStatsModule censusStats) {
186+
this.censusStatsOverride = censusStats;
203187
return thisT();
204188
}
205189

@@ -242,22 +226,11 @@ final List<ServerStreamTracer.Factory> getTracerFactories() {
242226
ArrayList<ServerStreamTracer.Factory> tracerFactories =
243227
new ArrayList<ServerStreamTracer.Factory>();
244228
if (statsEnabled) {
245-
Tagger tagger = this.tagger != null ? this.tagger : Tags.getTagger();
246-
TagContextBinarySerializer tagCtxSerializer =
247-
this.tagCtxSerializer != null
248-
? this.tagCtxSerializer
249-
: Tags.getTagPropagationComponent().getBinarySerializer();
250-
StatsRecorder statsRecorder =
251-
this.statsRecorder != null ? this.statsRecorder : Stats.getStatsRecorder();
252-
CensusStatsModule censusStats =
253-
new CensusStatsModule(
254-
tagger,
255-
tagCtxSerializer,
256-
statsRecorder,
257-
GrpcUtil.STOPWATCH_SUPPLIER,
258-
true,
259-
recordStats);
260-
tracerFactories.add(censusStats.getServerTracerFactory());
229+
CensusStatsModule censusStats = this.censusStatsOverride;
230+
if (censusStats == null) {
231+
censusStats = new CensusStatsModule(GrpcUtil.STOPWATCH_SUPPLIER, true);
232+
}
233+
tracerFactories.add(censusStats.getServerTracerFactory(recordStats));
261234
}
262235
if (tracingEnabled) {
263236
CensusTracingModule censusTracing =

‎core/src/main/java/io/grpc/internal/CensusStatsModule.java‎

Lines changed: 58 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,12 @@
3939
import io.grpc.StreamTracer;
4040
import io.opencensus.contrib.grpc.metrics.RpcMeasureConstants;
4141
import io.opencensus.stats.MeasureMap;
42+
import io.opencensus.stats.Stats;
4243
import io.opencensus.stats.StatsRecorder;
4344
import io.opencensus.tags.TagContext;
4445
import io.opencensus.tags.TagValue;
4546
import io.opencensus.tags.Tagger;
47+
import io.opencensus.tags.Tags;
4648
import io.opencensus.tags.propagation.TagContextBinarySerializer;
4749
import io.opencensus.tags.propagation.TagContextSerializationException;
4850
import java.util.concurrent.TimeUnit;
@@ -74,21 +76,33 @@ final class CensusStatsModule {
7476
private final Supplier<Stopwatch> stopwatchSupplier;
7577
@VisibleForTesting
7678
final Metadata.Key<TagContext> statsHeader;
77-
private final StatsClientInterceptor clientInterceptor = new StatsClientInterceptor();
78-
private final ServerTracerFactory serverTracerFactory = new ServerTracerFactory();
7979
private final boolean propagateTags;
80-
private final boolean recordStats;
8180

81+
/**
82+
* Creates a {@link CensusStatsModule} with the default OpenCensus implementation.
83+
*/
84+
CensusStatsModule(Supplier<Stopwatch> stopwatchSupplier, boolean propagateTags) {
85+
this(
86+
Tags.getTagger(),
87+
Tags.getTagPropagationComponent().getBinarySerializer(),
88+
Stats.getStatsRecorder(),
89+
stopwatchSupplier,
90+
propagateTags);
91+
}
92+
93+
/**
94+
* Creates a {@link CensusStatsModule} with the given OpenCensus implementation.
95+
*/
8296
CensusStatsModule(
8397
final Tagger tagger,
8498
final TagContextBinarySerializer tagCtxSerializer,
8599
StatsRecorder statsRecorder, Supplier<Stopwatch> stopwatchSupplier,
86-
boolean propagateTags, boolean recordStats) {
100+
boolean propagateTags) {
87101
this.tagger = checkNotNull(tagger, "tagger");
88102
this.statsRecorder = checkNotNull(statsRecorder, "statsRecorder");
103+
checkNotNull(tagCtxSerializer, "tagCtxSerializer");
89104
this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier");
90105
this.propagateTags = propagateTags;
91-
this.recordStats = recordStats;
92106
this.statsHeader =
93107
Metadata.Key.of("grpc-tags-bin", new Metadata.BinaryMarshaller<TagContext>() {
94108
@Override
@@ -118,22 +132,23 @@ public TagContext parseBytes(byte[] serialized) {
118132
* Creates a {@link ClientCallTracer} for a new call.
119133
*/
120134
@VisibleForTesting
121-
ClientCallTracer newClientCallTracer(TagContext parentCtx, String fullMethodName) {
122-
return new ClientCallTracer(this, parentCtx, fullMethodName);
135+
ClientCallTracer newClientCallTracer(
136+
TagContext parentCtx, String fullMethodName, boolean recordStats) {
137+
return new ClientCallTracer(this, parentCtx, fullMethodName, recordStats);
123138
}
124139

125140
/**
126141
* Returns the server tracer factory.
127142
*/
128-
ServerStreamTracer.Factory getServerTracerFactory() {
129-
return serverTracerFactory;
143+
ServerStreamTracer.Factory getServerTracerFactory(boolean recordStats) {
144+
return new ServerTracerFactory(recordStats);
130145
}
131146

132147
/**
133148
* Returns the client interceptor that facilitates Census-based stats reporting.
134149
*/
135-
ClientInterceptor getClientInterceptor() {
136-
return clientInterceptor;
150+
ClientInterceptor getClientInterceptor(boolean recordStats) {
151+
return new StatsClientInterceptor(recordStats);
137152
}
138153

139154
private static final class ClientTracer extends ClientStreamTracer {
@@ -206,12 +221,18 @@ static final class ClientCallTracer extends ClientStreamTracer.Factory {
206221
private volatile ClientTracer streamTracer;
207222
private volatile int callEnded;
208223
private final TagContext parentCtx;
224+
private final boolean recordStats;
209225

210-
ClientCallTracer(CensusStatsModule module, TagContext parentCtx, String fullMethodName) {
226+
ClientCallTracer(
227+
CensusStatsModule module,
228+
TagContext parentCtx,
229+
String fullMethodName,
230+
boolean recordStats) {
211231
this.module = module;
212232
this.parentCtx = checkNotNull(parentCtx, "parentCtx");
213233
this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
214234
this.stopwatch = module.stopwatchSupplier.get().start();
235+
this.recordStats = recordStats;
215236
}
216237

217238
@Override
@@ -241,7 +262,7 @@ void callEnded(Status status) {
241262
if (callEndedUpdater.getAndSet(this, 1) != 0) {
242263
return;
243264
}
244-
if (!module.recordStats) {
265+
if (!recordStats) {
245266
return;
246267
}
247268
stopwatch.stop();
@@ -299,6 +320,7 @@ private static final class ServerTracer extends ServerStreamTracer {
299320
private volatile int streamClosed;
300321
private final Stopwatch stopwatch;
301322
private final Tagger tagger;
323+
private final boolean recordStats;
302324
private volatile long outboundMessageCount;
303325
private volatile long inboundMessageCount;
304326
private volatile long outboundWireSize;
@@ -311,12 +333,14 @@ private static final class ServerTracer extends ServerStreamTracer {
311333
String fullMethodName,
312334
TagContext parentCtx,
313335
Supplier<Stopwatch> stopwatchSupplier,
314-
Tagger tagger) {
336+
Tagger tagger,
337+
boolean recordStats) {
315338
this.module = module;
316339
this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
317340
this.parentCtx = checkNotNull(parentCtx, "parentCtx");
318341
this.stopwatch = stopwatchSupplier.get().start();
319342
this.tagger = tagger;
343+
this.recordStats = recordStats;
320344
}
321345

322346
@Override
@@ -360,7 +384,7 @@ public void streamClosed(Status status) {
360384
if (streamClosedUpdater.getAndSet(this, 1) != 0) {
361385
return;
362386
}
363-
if (!module.recordStats) {
387+
if (!recordStats) {
364388
return;
365389
}
366390
stopwatch.stop();
@@ -397,6 +421,12 @@ public Context filterContext(Context context) {
397421

398422
@VisibleForTesting
399423
final class ServerTracerFactory extends ServerStreamTracer.Factory {
424+
private final boolean recordStats;
425+
426+
ServerTracerFactory(boolean recordStats) {
427+
this.recordStats = recordStats;
428+
}
429+
400430
@Override
401431
public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
402432
TagContext parentCtx = headers.get(statsHeader);
@@ -409,19 +439,30 @@ public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata
409439
.put(RpcMeasureConstants.RPC_METHOD, TagValue.create(fullMethodName))
410440
.build();
411441
return new ServerTracer(
412-
CensusStatsModule.this, fullMethodName, parentCtx, stopwatchSupplier, tagger);
442+
CensusStatsModule.this,
443+
fullMethodName,
444+
parentCtx,
445+
stopwatchSupplier,
446+
tagger,
447+
recordStats);
413448
}
414449
}
415450

416451
@VisibleForTesting
417452
final class StatsClientInterceptor implements ClientInterceptor {
453+
private final boolean recordStats;
454+
455+
StatsClientInterceptor(boolean recordStats) {
456+
this.recordStats = recordStats;
457+
}
458+
418459
@Override
419460
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
420461
MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
421462
// New RPCs on client-side inherit the tag context from the current Context.
422463
TagContext parentCtx = tagger.getCurrentTagContext();
423464
final ClientCallTracer tracerFactory =
424-
newClientCallTracer(parentCtx, method.getFullMethodName());
465+
newClientCallTracer(parentCtx, method.getFullMethodName(), recordStats);
425466
ClientCall<ReqT, RespT> call =
426467
next.newCall(method, callOptions.withStreamTracerFactory(tracerFactory));
427468
return new SimpleForwardingClientCall<ReqT, RespT>(call) {

‎core/src/test/java/io/grpc/internal/AbstractManagedChannelImplBuilderTest.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -393,12 +393,24 @@ public void idleTimeout() {
393393
static class Builder extends AbstractManagedChannelImplBuilder<Builder> {
394394
Builder(String target) {
395395
super(target);
396-
statsImplementation(DUMMY_TAGGER, DUMMY_TAG_CONTEXT_BINARY_SERIALIZER, DUMMY_STATS_RECORDER);
396+
overrideCensusStatsModule(
397+
new CensusStatsModule(
398+
DUMMY_TAGGER,
399+
DUMMY_TAG_CONTEXT_BINARY_SERIALIZER,
400+
DUMMY_STATS_RECORDER,
401+
GrpcUtil.STOPWATCH_SUPPLIER,
402+
true));
397403
}
398404

399405
Builder(SocketAddress directServerAddress, String authority) {
400406
super(directServerAddress, authority);
401-
statsImplementation(DUMMY_TAGGER, DUMMY_TAG_CONTEXT_BINARY_SERIALIZER, DUMMY_STATS_RECORDER);
407+
overrideCensusStatsModule(
408+
new CensusStatsModule(
409+
DUMMY_TAGGER,
410+
DUMMY_TAG_CONTEXT_BINARY_SERIALIZER,
411+
DUMMY_STATS_RECORDER,
412+
GrpcUtil.STOPWATCH_SUPPLIER,
413+
true));
402414
}
403415

404416
@Override

‎core/src/test/java/io/grpc/internal/AbstractServerImplBuilderTest.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,13 @@ public void getTracerFactories_disableBoth() {
145145

146146
static class Builder extends AbstractServerImplBuilder<Builder> {
147147
Builder() {
148-
statsImplementation(DUMMY_TAGGER, DUMMY_TAG_CONTEXT_BINARY_SERIALIZER, DUMMY_STATS_RECORDER);
148+
overrideCensusStatsModule(
149+
new CensusStatsModule(
150+
DUMMY_TAGGER,
151+
DUMMY_TAG_CONTEXT_BINARY_SERIALIZER,
152+
DUMMY_STATS_RECORDER,
153+
GrpcUtil.STOPWATCH_SUPPLIER,
154+
true));
149155
}
150156

151157
@Override

0 commit comments

Comments
 (0)