Skip to content

Commit

Permalink
Allow for more customisation of the tracing span #1303
Browse files Browse the repository at this point in the history 
Span customizer now accepts RedisCommand and Span to customize tracing spans based on the actual command.
  • Loading branch information
@mp911de
mp911de committed Jul 3, 2020
1 parent b4263f1 commit 0e14c54
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 45 deletions.
29 changes: 1 addition & 28 deletions src/main/java/io/lettuce/core/protocol/CommandHandler.java
View file
Expand Up @@ -90,8 +90,6 @@ public class CommandHandler extends ChannelDuplexHandler implements HasQueuedCom

private final boolean tracingEnabled;

private final boolean includeCommandArgsInSpanTags;

private final float discardReadBytesRatio;

private final boolean boundedQueues;
Expand Down Expand Up @@ -138,7 +136,6 @@ public CommandHandler(ClientOptions clientOptions, ClientResources clientResourc
Tracing tracing = clientResources.tracing();

this.tracingEnabled = tracing.isEnabled();
this.includeCommandArgsInSpanTags = tracing.includeCommandArgsInSpanTags();

float bufferUsageRatio = clientOptions.getBufferUsageRatio();
this.discardReadBytesRatio = bufferUsageRatio / (bufferUsageRatio + 1);
Expand Down Expand Up @@ -402,35 +399,11 @@ private void writeSingleCommand(ChannelHandlerContext ctx, RedisCommand<?, ?, ?>
TraceContext context = provider.getTraceContext();

Tracer.Span span = tracer.nextSpan(context);
span.name(command.getType().name());

if (includeCommandArgsInSpanTags && command.getArgs() != null) {
span.tag("redis.args", command.getArgs().toCommandString());
}

span.remoteEndpoint(tracedEndpoint);
span.start();
span.remoteEndpoint(tracedEndpoint).start(command);

if (traced != null) {
traced.setSpan(span);
}

CompleteableCommand<?> completeableCommand = (CompleteableCommand<?>) command;
completeableCommand.onComplete((o, throwable) -> {

if (command.getOutput() != null) {

String error = command.getOutput().getError();
if (error != null) {
span.tag("error", error);
} else if (throwable != null) {
span.tag("exception", throwable.toString());
span.error(throwable);
}
}

span.finish();
});
}

ctx.write(command, promise);
Expand Down
85 changes: 73 additions & 12 deletions src/main/java/io/lettuce/core/tracing/BraveTracing.java
View file
Expand Up @@ -17,12 +17,15 @@

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

import reactor.core.publisher.Mono;
import brave.Span;
import brave.propagation.TraceContextOrSamplingFlags;
import io.lettuce.core.internal.LettuceAssert;
import io.lettuce.core.protocol.CompleteableCommand;
import io.lettuce.core.protocol.RedisCommand;

/**
* {@link Tracing} integration with OpenZipkin's Brave {@link brave.Tracer}. This implementation creates Brave
Expand Down Expand Up @@ -71,7 +74,7 @@ private BraveTracing(Builder builder) {
LettuceAssert.notNull(builder.serviceName, "Service name must not be null");

this.tracingOptions = new BraveTracingOptions(builder.serviceName, builder.endpointCustomizer, builder.spanCustomizer);
this.tracer = new BraveTracer(builder.tracing, this.tracingOptions);
this.tracer = new BraveTracer(builder.tracing, this.tracingOptions, builder.includeCommandArgsInSpanTags);
this.includeCommandArgsInSpanTags = builder.includeCommandArgsInSpanTags;
}

Expand Down Expand Up @@ -109,7 +112,7 @@ public static class Builder {
private Consumer<zipkin2.Endpoint.Builder> endpointCustomizer = it -> {
};

private Consumer<brave.Span> spanCustomizer = it -> {
private BiConsumer<RedisCommand<Object, Object, Object>, Span> spanCustomizer = (command, span) -> {
};

private boolean includeCommandArgsInSpanTags = true;
Expand Down Expand Up @@ -193,6 +196,22 @@ public Builder spanCustomizer(Consumer<brave.Span> spanCustomizer) {

LettuceAssert.notNull(spanCustomizer, "Span customizer must not be null!");

this.spanCustomizer = (command, span) -> spanCustomizer.accept(span);
return this;
}

/**
* Sets an {@link brave.Span} customizer to customize the {@link brave.Span} based on the underlying
* {@link RedisCommand}. The customizer is invoked before {@link Span#finish()} finishing} the span.
*
* @param spanCustomizer must not be {@code null}.
* @return {@code this} {@link Builder}.
* @since 6.0
*/
public Builder spanCustomizer(BiConsumer<RedisCommand<Object, Object, Object>, brave.Span> spanCustomizer) {

LettuceAssert.notNull(spanCustomizer, "Span customizer must not be null!");

this.spanCustomizer = spanCustomizer;
return this;
}
Expand Down Expand Up @@ -257,9 +276,12 @@ static class BraveTracer extends Tracer {

private final BraveTracingOptions tracingOptions;

BraveTracer(brave.Tracing tracing, BraveTracingOptions tracingOptions) {
private final boolean includeCommandArgsInSpanTags;

BraveTracer(brave.Tracing tracing, BraveTracingOptions tracingOptions, boolean includeCommandArgsInSpanTags) {
this.tracing = tracing;
this.tracingOptions = tracingOptions;
this.includeCommandArgsInSpanTags = includeCommandArgsInSpanTags;
}

@Override
Expand Down Expand Up @@ -290,7 +312,7 @@ private Span postProcessSpan(brave.Span span) {
return NoOpTracing.NoOpSpan.INSTANCE;
}

return new BraveSpan(span.kind(brave.Span.Kind.CLIENT), this.tracingOptions);
return new BraveSpan(span.kind(brave.Span.Kind.CLIENT), this.tracingOptions, includeCommandArgsInSpanTags);
}

}
Expand All @@ -304,15 +326,47 @@ static class BraveSpan extends Tracer.Span {

private final BraveTracingOptions tracingOptions;

BraveSpan(Span span, BraveTracingOptions tracingOptions) {
private final boolean includeCommandArgsInSpanTags;

BraveSpan(Span span, BraveTracingOptions tracingOptions, boolean includeCommandArgsInSpanTags) {
this.span = span;
this.tracingOptions = tracingOptions;
this.includeCommandArgsInSpanTags = includeCommandArgsInSpanTags;
}

@Override
public BraveSpan start() {
public BraveSpan start(RedisCommand<?, ?, ?> command) {

span.name(command.getType().name());

if (includeCommandArgsInSpanTags && command.getArgs() != null) {
span.tag("redis.args", command.getArgs().toCommandString());
}

if (command instanceof CompleteableCommand) {
CompleteableCommand<?> completeableCommand = (CompleteableCommand<?>) command;
completeableCommand.onComplete((o, throwable) -> {

if (command.getOutput() != null) {

String error = command.getOutput().getError();
if (error != null) {
span.tag("error", error);
} else if (throwable != null) {
span.tag("exception", throwable.toString());
span.error(throwable);
}
}

span.finish();
});
} else {
throw new IllegalArgumentException("Command " + command
+ " must implement CompleteableCommand to attach Span completion to command completion");
}

span.start();
this.tracingOptions.customizeSpan(command, span);

return this;
}
Expand Down Expand Up @@ -352,15 +406,21 @@ public BraveSpan error(Throwable throwable) {
@Override
public BraveSpan remoteEndpoint(Endpoint endpoint) {

span.remoteEndpoint(BraveEndpoint.class.cast(endpoint).endpoint);
zipkin2.Endpoint zkEndpoint = BraveEndpoint.class.cast(endpoint).endpoint;

if (zkEndpoint.serviceName() != null) {
span.remoteServiceName(zkEndpoint.serviceName());
}

String ip = zkEndpoint.ipv6() != null ? zkEndpoint.ipv6() : zkEndpoint.ipv4();
span.remoteIpAndPort(ip, zkEndpoint.portAsInt());

return this;
}

@Override
public void finish() {

this.tracingOptions.customizeSpan(span);
span.finish();
}

Expand Down Expand Up @@ -448,10 +508,10 @@ static class BraveTracingOptions {

private final Consumer<zipkin2.Endpoint.Builder> endpointCustomizer;

private final Consumer<brave.Span> spanCustomizer;
private final BiConsumer<RedisCommand<Object, Object, Object>, brave.Span> spanCustomizer;

BraveTracingOptions(String serviceName, Consumer<zipkin2.Endpoint.Builder> endpointCustomizer,
Consumer<Span> spanCustomizer) {
BiConsumer<RedisCommand<Object, Object, Object>, brave.Span> spanCustomizer) {
this.serviceName = serviceName;
this.endpointCustomizer = endpointCustomizer;
this.spanCustomizer = spanCustomizer;
Expand All @@ -461,8 +521,9 @@ void customizeEndpoint(zipkin2.Endpoint.Builder builder) {
this.endpointCustomizer.accept(builder);
}

void customizeSpan(brave.Span span) {
this.spanCustomizer.accept(span);
@SuppressWarnings("unchecked")
void customizeSpan(RedisCommand<?, ?, ?> command, brave.Span span) {
this.spanCustomizer.accept((RedisCommand<Object, Object, Object>) command, span);
}

}
Expand Down
4 changes: 3 additions & 1 deletion src/main/java/io/lettuce/core/tracing/NoOpTracing.java
View file
Expand Up @@ -15,6 +15,8 @@
*/
package io.lettuce.core.tracing;

import io.lettuce.core.protocol.RedisCommand;

import java.net.SocketAddress;

/**
Expand Down Expand Up @@ -87,7 +89,7 @@ public static class NoOpSpan extends Tracer.Span {
static final NoOpSpan INSTANCE = new NoOpSpan();

@Override
public Tracer.Span start() {
public Tracer.Span start(RedisCommand<?, ?, ?> command) {
return this;
}

Expand Down
5 changes: 4 additions & 1 deletion src/main/java/io/lettuce/core/tracing/Tracer.java
View file
Expand Up @@ -15,6 +15,8 @@
*/
package io.lettuce.core.tracing;

import io.lettuce.core.protocol.RedisCommand;

/**
* Tracing abstraction to create {@link Span}s to capture latency and behavior of Redis commands.
*
Expand Down Expand Up @@ -48,9 +50,10 @@ public abstract static class Span {
/**
* Starts the span with.
*
* @param command the underlying command.
* @return {@literal this} {@link Span}.
*/
public abstract Span start();
public abstract Span start(RedisCommand<?, ?, ?> command);

/**
* Sets the name for this {@link Span}.
Expand Down
14 changes: 11 additions & 3 deletions src/test/java/io/lettuce/core/tracing/BraveTracingUnitTests.java
View file
Expand Up @@ -17,6 +17,7 @@

import static org.assertj.core.api.Assertions.assertThat;

import java.util.List;
import java.util.Queue;
import java.util.concurrent.LinkedBlockingQueue;

Expand All @@ -26,15 +27,21 @@
import org.junit.jupiter.api.Test;
import org.springframework.test.util.ReflectionTestUtils;

import io.lettuce.core.protocol.AsyncCommand;
import zipkin2.Span;
import brave.Tag;
import brave.Tracer;
import brave.Tracing;
import brave.handler.MutableSpan;
import brave.propagation.CurrentTraceContext;
import io.lettuce.core.TestSupport;
import io.lettuce.core.protocol.Command;
import io.lettuce.core.protocol.CommandType;
import io.netty.channel.unix.DomainSocketAddress;

/**
* Unit tests for {@link BraveTracing}.
*
* @author Mark Paluch
* @author Daniel Albuquerque
*/
Expand Down Expand Up @@ -109,13 +116,14 @@ void shouldCustomizeEndpoint() {
void shouldCustomizeSpan() {

BraveTracing tracing = BraveTracing.builder().tracing(clientTracing)
.spanCustomizer(it -> it.remoteServiceName("remote")).build();
.spanCustomizer((command, span) -> span.tag("cmd", command.getType().name())).build();

BraveTracing.BraveSpan span = (BraveTracing.BraveSpan) tracing.getTracerProvider().getTracer().nextSpan();
span.finish();
span.start(new AsyncCommand<>(new Command<>(CommandType.AUTH, null)));

MutableSpan braveSpan = (MutableSpan) ReflectionTestUtils.getField(span.getSpan(), "state");
List<String> tags = (List) ReflectionTestUtils.getField(braveSpan, "tags");

assertThat(braveSpan.remoteServiceName()).isEqualTo("remote");
assertThat(tags).contains("cmd", "AUTH");
}
}

0 comments on commit 0e14c54

Please sign in to comment.