Skip to content

Commit

Permalink
feat: fix review changes
Browse files Browse the repository at this point in the history
  • Loading branch information
suraj-qlogic committed Jun 4, 2020
1 parent c6806ab commit 9d85b03
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 42 deletions.
Expand Up @@ -136,41 +136,35 @@ public T call() throws DatastoreException {

@Override
public <T> T runInTransaction(final TransactionCallable<T> callable) {
final DatastoreImpl self = this;
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new ReadWriteTransactionCallable<T>(self, callable, null),
new ReadWriteTransactionCallable<T>(this, callable, null),
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}

@Override
public <T> T runInTransaction(
final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
final DatastoreImpl self = this;
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new ReadWriteTransactionCallable<T>(self, callable, transactionOptions),
new ReadWriteTransactionCallable<T>(this, callable, transactionOptions),
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand All @@ -192,8 +186,7 @@ <T> QueryResults<T> run(com.google.datastore.v1.ReadOptions readOptionsPb, Query
com.google.datastore.v1.RunQueryResponse runQuery(
final com.google.datastore.v1.RunQueryRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.RunQueryResponse>() {
@Override
Expand All @@ -210,7 +203,6 @@ public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand Down Expand Up @@ -252,8 +244,7 @@ public List<Key> allocateId(IncompleteKey... keys) {
private com.google.datastore.v1.AllocateIdsResponse allocateIds(
final com.google.datastore.v1.AllocateIdsRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.AllocateIdsResponse>() {
@Override
Expand All @@ -268,7 +259,6 @@ public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreExcept
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand Down Expand Up @@ -418,8 +408,7 @@ protected Entity computeNext() {
com.google.datastore.v1.LookupResponse lookup(
final com.google.datastore.v1.LookupRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.LookupResponse>() {
@Override
Expand All @@ -436,7 +425,6 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand All @@ -460,8 +448,7 @@ public List<Key> reserveIds(Key... keys) {
com.google.datastore.v1.ReserveIdsResponse reserveIds(
final com.google.datastore.v1.ReserveIdsRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
@Override
Expand All @@ -476,7 +463,6 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand Down Expand Up @@ -570,8 +556,7 @@ private com.google.datastore.v1.CommitResponse commitMutation(
com.google.datastore.v1.CommitResponse commit(
final com.google.datastore.v1.CommitRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.CommitResponse>() {
@Override
Expand All @@ -588,7 +573,6 @@ public com.google.datastore.v1.CommitResponse call() throws DatastoreException {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand All @@ -601,8 +585,7 @@ ByteString requestTransactionId(
com.google.datastore.v1.BeginTransactionResponse beginTransaction(
final com.google.datastore.v1.BeginTransactionRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
return RetryHelper.runWithRetries(
new Callable<com.google.datastore.v1.BeginTransactionResponse>() {
@Override
Expand All @@ -618,7 +601,6 @@ public com.google.datastore.v1.BeginTransactionResponse call()
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand All @@ -632,8 +614,7 @@ void rollbackTransaction(ByteString transaction) {

void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
Scope scope = traceUtil.getTracer().withSpan(span);
try {
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
RetryHelper.runWithRetries(
new Callable<Void>() {
@Override
Expand All @@ -649,7 +630,6 @@ public Void call() throws DatastoreException {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
throw DatastoreException.translateAndThrow(e);
} finally {
scope.close();
span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
Expand Down
Expand Up @@ -23,14 +23,14 @@
import io.opencensus.trace.Tracing;

/**
* Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with Open
* Census APIs.
* Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with
* OpenCensus APIs.
*
* <p>TraceUtil Instance are created by the {@link TraceUtil#getInstance()} method.
* <p>TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
*/
public class TraceUtil {
private final Tracer tracer = Tracing.getTracer();
private static TraceUtil traceUtil;
private static final TraceUtil traceUtil = new TraceUtil();
static final String SPAN_NAME_ALLOCATEIDS = "CloudDatastoreOperation.allocateIds";
static final String SPAN_NAME_TRANSACTION = "CloudDatastoreOperation.readWriteTransaction";
static final String SPAN_NAME_BEGINTRANSACTION = "CloudDatastoreOperation.beginTransaction";
Expand Down Expand Up @@ -67,9 +67,6 @@ public Tracer getTracer() {
* @return An instance of {@link TraceUtil}
*/
public static TraceUtil getInstance() {
if (traceUtil == null) {
traceUtil = new TraceUtil();
}
return traceUtil;
}

Expand Down
Expand Up @@ -77,18 +77,18 @@ public HttpDatastoreRpc(DatastoreOptions options) {

private HttpRequestInitializer getHttpRequestInitializer(
final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
HttpRequestInitializer delegate = httpTransportOptions.getHttpRequestInitializer(options);
// Open Census initialization
CensusHttpModule censusHttpModule =
new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
delegate = censusHttpModule.getHttpRequestInitializer(delegate);
final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
censusHttpModule.getHttpRequestInitializer(
httpTransportOptions.getHttpRequestInitializer(options));

final String applicationName = options.getApplicationName();
final HttpRequestInitializer localDelegate = delegate;
return new HttpRequestInitializer() {
@Override
public void initialize(HttpRequest httpRequest) throws IOException {
localDelegate.initialize(httpRequest);
censusHttpModuleHttpRequestInitializer.initialize(httpRequest);
httpRequest.getHeaders().setUserAgent(applicationName);
}
};
Expand Down

0 comments on commit 9d85b03

Please sign in to comment.