1 /* 2 * Copyright 2016 The gRPC Authors 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package io.grpc.internal.testing; 18 19 import static com.google.common.base.Charsets.UTF_8; 20 import static com.google.common.base.Preconditions.checkNotNull; 21 22 import com.google.common.base.Function; 23 import com.google.common.collect.ImmutableMap; 24 import com.google.common.collect.Iterators; 25 import com.google.common.collect.Maps; 26 import io.opencensus.common.Scope; 27 import io.opencensus.stats.Measure; 28 import io.opencensus.stats.MeasureMap; 29 import io.opencensus.stats.StatsRecorder; 30 import io.opencensus.tags.Tag; 31 import io.opencensus.tags.TagContext; 32 import io.opencensus.tags.TagContextBuilder; 33 import io.opencensus.tags.TagKey; 34 import io.opencensus.tags.TagValue; 35 import io.opencensus.tags.Tagger; 36 import io.opencensus.tags.propagation.TagContextBinarySerializer; 37 import io.opencensus.tags.propagation.TagContextDeserializationException; 38 import io.opencensus.tags.unsafe.ContextUtils; 39 import io.opencensus.trace.Annotation; 40 import io.opencensus.trace.AttributeValue; 41 import io.opencensus.trace.EndSpanOptions; 42 import io.opencensus.trace.Link; 43 import io.opencensus.trace.MessageEvent; 44 import io.opencensus.trace.Sampler; 45 import io.opencensus.trace.Span; 46 import io.opencensus.trace.SpanBuilder; 47 import io.opencensus.trace.SpanContext; 48 import io.opencensus.trace.SpanId; 49 import io.opencensus.trace.TraceId; 50 import io.opencensus.trace.TraceOptions; 51 import java.util.EnumSet; 52 import java.util.Iterator; 53 import java.util.List; 54 import java.util.Map; 55 import java.util.Random; 56 import java.util.concurrent.BlockingQueue; 57 import java.util.concurrent.LinkedBlockingQueue; 58 import java.util.concurrent.TimeUnit; 59 import javax.annotation.Nullable; 60 61 public class StatsTestUtils { 62 private StatsTestUtils() { 63 } 64 65 public static class MetricsRecord { 66 67 public final ImmutableMap<TagKey, TagValue> tags; 68 public final ImmutableMap<Measure, Number> metrics; 69 70 private MetricsRecord( 71 ImmutableMap<TagKey, TagValue> tags, ImmutableMap<Measure, Number> metrics) { 72 this.tags = tags; 73 this.metrics = metrics; 74 } 75 76 /** 77 * Returns the value of a metric, or {@code null} if not found. 78 */ 79 @Nullable 80 public Double getMetric(Measure measure) { 81 for (Map.Entry<Measure, Number> m : metrics.entrySet()) { 82 if (m.getKey().equals(measure)) { 83 Number value = m.getValue(); 84 if (value instanceof Double) { 85 return (Double) value; 86 } else if (value instanceof Long) { 87 return (double) (Long) value; 88 } 89 throw new AssertionError("Unexpected measure value type: " + value.getClass().getName()); 90 } 91 } 92 return null; 93 } 94 95 /** 96 * Returns the value of a metric converted to long, or throw if not found. 97 */ 98 public long getMetricAsLongOrFail(Measure measure) { 99 Double doubleValue = getMetric(measure); 100 checkNotNull(doubleValue, "Measure not found: %s", measure.getName()); 101 long longValue = (long) (Math.abs(doubleValue) + 0.0001); 102 if (doubleValue < 0) { 103 longValue = -longValue; 104 } 105 return longValue; 106 } 107 108 @Override 109 public String toString() { 110 return "[tags=" + tags + ", metrics=" + metrics + "]"; 111 } 112 } 113 114 /** 115 * This tag will be propagated by {@link FakeTagger} on the wire. 116 */ 117 public static final TagKey EXTRA_TAG = TagKey.create("/rpc/test/extratag"); 118 119 private static final String EXTRA_TAG_HEADER_VALUE_PREFIX = "extratag:"; 120 121 /** 122 * A {@link Tagger} implementation that saves metrics records to be accessible from {@link 123 * #pollRecord()} and {@link #pollRecord(long, TimeUnit)}, until {@link #rolloverRecords} is 124 * called. 125 */ 126 public static final class FakeStatsRecorder extends StatsRecorder { 127 128 private BlockingQueue<MetricsRecord> records; 129 130 public FakeStatsRecorder() { 131 rolloverRecords(); 132 } 133 134 @Override 135 public MeasureMap newMeasureMap() { 136 return new FakeStatsRecord(this); 137 } 138 139 public MetricsRecord pollRecord() { 140 return getCurrentRecordSink().poll(); 141 } 142 143 public MetricsRecord pollRecord(long timeout, TimeUnit unit) throws InterruptedException { 144 return getCurrentRecordSink().poll(timeout, unit); 145 } 146 147 /** 148 * Disconnect this tagger with the contexts it has created so far. The records from those 149 * contexts will not show up in {@link #pollRecord}. Useful for isolating the records between 150 * test cases. 151 */ 152 // This needs to be synchronized with getCurrentRecordSink() which may run concurrently. 153 public synchronized void rolloverRecords() { 154 records = new LinkedBlockingQueue<MetricsRecord>(); 155 } 156 157 private synchronized BlockingQueue<MetricsRecord> getCurrentRecordSink() { 158 return records; 159 } 160 } 161 162 public static final class FakeTagger extends Tagger { 163 164 @Override 165 public FakeTagContext empty() { 166 return FakeTagContext.EMPTY; 167 } 168 169 @Override 170 public TagContext getCurrentTagContext() { 171 return ContextUtils.TAG_CONTEXT_KEY.get(); 172 } 173 174 @Override 175 public TagContextBuilder emptyBuilder() { 176 return new FakeTagContextBuilder(ImmutableMap.<TagKey, TagValue>of()); 177 } 178 179 @Override 180 public FakeTagContextBuilder toBuilder(TagContext tags) { 181 return new FakeTagContextBuilder(getTags(tags)); 182 } 183 184 @Override 185 public TagContextBuilder currentBuilder() { 186 throw new UnsupportedOperationException(); 187 } 188 189 @Override 190 public Scope withTagContext(TagContext tags) { 191 throw new UnsupportedOperationException(); 192 } 193 } 194 195 public static final class FakeTagContextBinarySerializer extends TagContextBinarySerializer { 196 197 private final FakeTagger tagger = new FakeTagger(); 198 199 @Override 200 public TagContext fromByteArray(byte[] bytes) throws TagContextDeserializationException { 201 String serializedString = new String(bytes, UTF_8); 202 if (serializedString.startsWith(EXTRA_TAG_HEADER_VALUE_PREFIX)) { 203 return tagger.emptyBuilder() 204 .put(EXTRA_TAG, 205 TagValue.create(serializedString.substring(EXTRA_TAG_HEADER_VALUE_PREFIX.length()))) 206 .build(); 207 } else { 208 throw new TagContextDeserializationException("Malformed value"); 209 } 210 } 211 212 @Override 213 public byte[] toByteArray(TagContext tags) { 214 TagValue extraTagValue = getTags(tags).get(EXTRA_TAG); 215 if (extraTagValue == null) { 216 throw new UnsupportedOperationException("TagContext must contain EXTRA_TAG"); 217 } 218 return (EXTRA_TAG_HEADER_VALUE_PREFIX + extraTagValue.asString()).getBytes(UTF_8); 219 } 220 } 221 222 public static final class FakeStatsRecord extends MeasureMap { 223 224 private final BlockingQueue<MetricsRecord> recordSink; 225 public final Map<Measure, Number> metrics = Maps.newHashMap(); 226 227 private FakeStatsRecord(FakeStatsRecorder statsRecorder) { 228 this.recordSink = statsRecorder.getCurrentRecordSink(); 229 } 230 231 @Override 232 public MeasureMap put(Measure.MeasureDouble measure, double value) { 233 metrics.put(measure, value); 234 return this; 235 } 236 237 @Override 238 public MeasureMap put(Measure.MeasureLong measure, long value) { 239 metrics.put(measure, value); 240 return this; 241 } 242 243 @Override 244 public void record(TagContext tags) { 245 recordSink.add(new MetricsRecord(getTags(tags), ImmutableMap.copyOf(metrics))); 246 } 247 248 @Override 249 public void record() { 250 throw new UnsupportedOperationException(); 251 } 252 } 253 254 public static final class FakeTagContext extends TagContext { 255 256 private static final FakeTagContext EMPTY = 257 new FakeTagContext(ImmutableMap.<TagKey, TagValue>of()); 258 259 private final ImmutableMap<TagKey, TagValue> tags; 260 261 private FakeTagContext(ImmutableMap<TagKey, TagValue> tags) { 262 this.tags = tags; 263 } 264 265 public ImmutableMap<TagKey, TagValue> getTags() { 266 return tags; 267 } 268 269 @Override 270 public String toString() { 271 return "[tags=" + tags + "]"; 272 } 273 274 @Override 275 protected Iterator<Tag> getIterator() { 276 return Iterators.transform( 277 tags.entrySet().iterator(), 278 new Function<Map.Entry<TagKey, TagValue>, Tag>() { 279 @Override 280 public Tag apply(@Nullable Map.Entry<TagKey, TagValue> entry) { 281 return Tag.create(entry.getKey(), entry.getValue()); 282 } 283 }); 284 } 285 } 286 287 public static class FakeTagContextBuilder extends TagContextBuilder { 288 289 private final Map<TagKey, TagValue> tagsBuilder = Maps.newHashMap(); 290 291 private FakeTagContextBuilder(Map<TagKey, TagValue> tags) { 292 tagsBuilder.putAll(tags); 293 } 294 295 @Override 296 public TagContextBuilder put(TagKey key, TagValue value) { 297 tagsBuilder.put(key, value); 298 return this; 299 } 300 301 @Override 302 public TagContextBuilder remove(TagKey key) { 303 tagsBuilder.remove(key); 304 return this; 305 } 306 307 @Override 308 public TagContext build() { 309 FakeTagContext context = new FakeTagContext(ImmutableMap.copyOf(tagsBuilder)); 310 return context; 311 } 312 313 @Override 314 public Scope buildScoped() { 315 throw new UnsupportedOperationException(); 316 } 317 } 318 319 // This method handles the default TagContext, which isn't an instance of FakeTagContext. 320 private static ImmutableMap<TagKey, TagValue> getTags(TagContext tags) { 321 return tags instanceof FakeTagContext 322 ? ((FakeTagContext) tags).getTags() 323 : ImmutableMap.<TagKey, TagValue>of(); 324 } 325 326 // TODO(bdrutu): Remove this class after OpenCensus releases support for this class. 327 public static class MockableSpan extends Span { 328 /** 329 * Creates a MockableSpan with a random trace ID and span ID. 330 */ 331 public static MockableSpan generateRandomSpan(Random random) { 332 return new MockableSpan( 333 SpanContext.create( 334 TraceId.generateRandomId(random), 335 SpanId.generateRandomId(random), 336 TraceOptions.DEFAULT), 337 null); 338 } 339 340 @Override 341 public void putAttributes(Map<String, AttributeValue> attributes) {} 342 343 @Override 344 public void addAnnotation(String description, Map<String, AttributeValue> attributes) {} 345 346 @Override 347 public void addAnnotation(Annotation annotation) {} 348 349 @Override 350 public void addMessageEvent(MessageEvent messageEvent) {} 351 352 @Override 353 public void addLink(Link link) {} 354 355 @Override 356 public void end(EndSpanOptions options) {} 357 358 private MockableSpan(SpanContext context, @Nullable EnumSet<Options> options) { 359 super(context, options); 360 } 361 362 /** 363 * Mockable implementation for the {@link SpanBuilder} class. 364 * 365 * <p>Not {@code final} to allow easy mocking. 366 * 367 */ 368 public static class Builder extends SpanBuilder { 369 370 @Override 371 public SpanBuilder setSampler(Sampler sampler) { 372 return this; 373 } 374 375 @Override 376 public SpanBuilder setParentLinks(List<Span> parentLinks) { 377 return this; 378 } 379 380 @Override 381 public SpanBuilder setRecordEvents(boolean recordEvents) { 382 return this; 383 } 384 385 @Override 386 public Span startSpan() { 387 return null; 388 } 389 } 390 } 391 } 392
