1717package com .google .cloud .spanner ;
1818
1919import com .google .cloud .Timestamp ;
20+
2021import com .google .common .util .concurrent .ListenableFuture ;
2122
22- class DatabaseClientImpl implements DatabaseClient {
23+ import io .opencensus .common .Scope ;
24+ import io .opencensus .trace .Span ;
25+ import io .opencensus .trace .Tracer ;
26+ import io .opencensus .trace .Tracing ;
2327
28+ class DatabaseClientImpl implements DatabaseClient {
29+ private static final String READ_WRITE_TRANSACTION = "CloudSpanner.ReadWriteTransaction" ;
30+ private static final String READ_ONLY_TRANSACTION = "CloudSpanner.ReadOnlyTransaction" ;
31+ private static final Tracer tracer = Tracing .getTracer ();
32+
33+ static {
34+ TraceUtil .exportSpans (READ_WRITE_TRANSACTION , READ_ONLY_TRANSACTION );
35+ }
36+
2437 private final SessionPool pool ;
2538
2639 DatabaseClientImpl (SessionPool pool ) {
@@ -29,47 +42,106 @@ class DatabaseClientImpl implements DatabaseClient {
2942
3043 @ Override
3144 public Timestamp write (Iterable <Mutation > mutations ) throws SpannerException {
32- return pool .getReadWriteSession ().write (mutations );
45+ Span span = tracer .spanBuilder (READ_WRITE_TRANSACTION ).startSpan ();
46+ try (Scope s = tracer .withSpan (span )) {
47+ return pool .getReadWriteSession ().write (mutations );
48+ } catch (RuntimeException e ) {
49+ TraceUtil .endSpanWithFailure (span , e );
50+ throw e ;
51+ } finally {
52+ span .end ();
53+ }
3354 }
3455
3556 @ Override
3657 public Timestamp writeAtLeastOnce (Iterable <Mutation > mutations ) throws SpannerException {
37- return pool .getReadSession ().writeAtLeastOnce (mutations );
58+ Span span = tracer .spanBuilder (READ_WRITE_TRANSACTION ).startSpan ();
59+ try (Scope s = tracer .withSpan (span )) {
60+ return pool .getReadSession ().writeAtLeastOnce (mutations );
61+ } catch (RuntimeException e ) {
62+ TraceUtil .endSpanWithFailure (span , e );
63+ throw e ;
64+ } finally {
65+ span .end ();
66+ }
3867 }
3968
4069 @ Override
4170 public ReadContext singleUse () {
42- return pool .getReadSession ().singleUse ();
71+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
72+ try (Scope s = tracer .withSpan (span )) {
73+ return pool .getReadSession ().singleUse ();
74+ } catch (RuntimeException e ) {
75+ TraceUtil .endSpanWithFailure (span , e );
76+ throw e ;
77+ }
78+
4379 }
4480
4581 @ Override
4682 public ReadContext singleUse (TimestampBound bound ) {
47- return pool .getReadSession ().singleUse (bound );
83+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
84+ try (Scope s = tracer .withSpan (span )) {
85+ return pool .getReadSession ().singleUse (bound );
86+ } catch (RuntimeException e ) {
87+ TraceUtil .endSpanWithFailure (span , e );
88+ throw e ;
89+ }
4890 }
4991
5092 @ Override
5193 public ReadOnlyTransaction singleUseReadOnlyTransaction () {
52- return pool .getReadSession ().singleUseReadOnlyTransaction ();
94+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
95+ try (Scope s = tracer .withSpan (span )) {
96+ return pool .getReadSession ().singleUseReadOnlyTransaction ();
97+ } catch (RuntimeException e ) {
98+ TraceUtil .endSpanWithFailure (span , e );
99+ throw e ;
100+ }
53101 }
54102
55103 @ Override
56104 public ReadOnlyTransaction singleUseReadOnlyTransaction (TimestampBound bound ) {
57- return pool .getReadSession ().singleUseReadOnlyTransaction (bound );
105+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
106+ try (Scope s = tracer .withSpan (span )) {
107+ return pool .getReadSession ().singleUseReadOnlyTransaction (bound );
108+ } catch (RuntimeException e ) {
109+ TraceUtil .endSpanWithFailure (span , e );
110+ throw e ;
111+ }
58112 }
59113
60114 @ Override
61115 public ReadOnlyTransaction readOnlyTransaction () {
62- return pool .getReadSession ().readOnlyTransaction ();
116+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
117+ try (Scope s = tracer .withSpan (span )) {
118+ return pool .getReadSession ().readOnlyTransaction ();
119+ } catch (RuntimeException e ) {
120+ TraceUtil .endSpanWithFailure (span , e );
121+ throw e ;
122+ }
63123 }
64124
65125 @ Override
66126 public ReadOnlyTransaction readOnlyTransaction (TimestampBound bound ) {
67- return pool .getReadSession ().readOnlyTransaction (bound );
127+ Span span = tracer .spanBuilder (READ_ONLY_TRANSACTION ).startSpan ();
128+ try (Scope s = tracer .withSpan (span )) {
129+ return pool .getReadSession ().readOnlyTransaction (bound );
130+ } catch (RuntimeException e ) {
131+ TraceUtil .endSpanWithFailure (span , e );
132+ throw e ;
133+ }
68134 }
69135
70136 @ Override
71137 public TransactionRunner readWriteTransaction () {
72- return pool .getReadWriteSession ().readWriteTransaction ();
138+ Span span = tracer .spanBuilder (READ_WRITE_TRANSACTION ).startSpan ();
139+ try (Scope s = tracer .withSpan (span )) {
140+ return pool .getReadWriteSession ().readWriteTransaction ();
141+ } catch (RuntimeException e ) {
142+ TraceUtil .endSpanWithFailure (span , e );
143+ throw e ;
144+ }
73145 }
74146
75147 ListenableFuture <Void > closeAsync () {
0 commit comments