Stackdriver Trace exporter: allow override of trace_service_stub. (#341)

Add Register(rvalue) and mark Register(ref) deprecated.

Author: g-easy

opencensus/exporters/trace/stackdriver/internal/stackdriver_exporter.cc

@@ -239,18 +239,13 @@ void ConvertSpans(
 
 class Handler : public ::opencensus::trace::exporter::SpanExporter::Handler {
  public:
-  Handler(const StackdriverOptions& opts,
-          const std::shared_ptr<grpc::Channel>& channel)
-      : opts_(opts),
-        stub_(::google::devtools::cloudtrace::v2::TraceService::NewStub(
-            channel)) {}
+  Handler(StackdriverOptions&& opts) : opts_(std::move(opts)) {}
 
   void Export(const std::vector<::opencensus::trace::exporter::SpanData>& spans)
       override;
 
  private:
   const StackdriverOptions opts_;
-  std::unique_ptr<google::devtools::cloudtrace::v2::TraceService::Stub> stub_;
 };
 
 void Handler::Export(
@@ -261,29 +256,52 @@ void Handler::Export(
   ::google::protobuf::Empty response;
   grpc::ClientContext context;
   context.set_deadline(absl::ToChronoTime(absl::Now() + opts_.rpc_deadline));
-  grpc::Status status = stub_->BatchWriteSpans(&context, request, &response);
+  grpc::Status status =
+      opts_.trace_service_stub->BatchWriteSpans(&context, request, &response);
   if (!status.ok()) {
     std::cerr << "BatchWriteSpans failed: "
               << opencensus::common::ToString(status) << "\n";
   }
 }
 
-}  // namespace
-
-// static
-void StackdriverExporter::Register(const StackdriverOptions& opts) {
+std::unique_ptr<google::devtools::cloudtrace::v2::TraceService::Stub>
+MakeStackdriverStub() {
   auto channel = ::grpc::CreateCustomChannel(
       kGoogleStackdriverTraceAddress, ::grpc::GoogleDefaultCredentials(),
       ::opencensus::common::WithUserAgent());
+  return ::google::devtools::cloudtrace::v2::TraceService::NewStub(channel);
+}
+
+}  // namespace
+
+// static
+void StackdriverExporter::Register(StackdriverOptions&& opts) {
+  if (opts.trace_service_stub == nullptr) {
+    opts.trace_service_stub = MakeStackdriverStub();
+  }
+  ::opencensus::trace::exporter::SpanExporter::RegisterHandler(
+      absl::make_unique<Handler>(std::move(opts)));
+}
+
+// static, DEPRECATED
+void StackdriverExporter::Register(StackdriverOptions& opts) {
+  if (opts.trace_service_stub == nullptr) {
+    opts.trace_service_stub = MakeStackdriverStub();
+  }
+  // Copy opts but take ownership of trace_service_stub.
+  StackdriverOptions copied_opts;
+  copied_opts.project_id = opts.project_id;
+  copied_opts.rpc_deadline = opts.rpc_deadline;
+  copied_opts.trace_service_stub = std::move(opts.trace_service_stub);
   ::opencensus::trace::exporter::SpanExporter::RegisterHandler(
-      absl::make_unique<Handler>(opts, channel));
+      absl::make_unique<Handler>(std::move(copied_opts)));
 }
 
 // static, DEPRECATED
 void StackdriverExporter::Register(absl::string_view project_id) {
   StackdriverOptions opts;
   opts.project_id = std::string(project_id);
-  Register(opts);
+  Register(std::move(opts));
 }
 
 }  // namespace trace

opencensus/exporters/trace/stackdriver/stackdriver_exporter.h

@@ -15,11 +15,13 @@
 #ifndef OPENCENSUS_EXPORTERS_TRACE_STACKDRIVER_STACKDRIVER_EXPORTER_H_
 #define OPENCENSUS_EXPORTERS_TRACE_STACKDRIVER_STACKDRIVER_EXPORTER_H_
 
+#include <memory>
 #include <string>
 
 #include "absl/base/macros.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
+#include "google/devtools/cloudtrace/v2/tracing.grpc.pb.h"
 
 namespace opencensus {
 namespace exporters {
@@ -31,12 +33,25 @@ struct StackdriverOptions {
 
   // The RPC deadline to use when exporting to Stackdriver.
   absl::Duration rpc_deadline = absl::Seconds(5);
+
+  // (optional) By default, the exporter connects to Stackdriver using gRPC. If
+  // this stub is non-null, the exporter will use this stub to send gRPC calls
+  // instead. Useful for testing.
+  std::unique_ptr<google::devtools::cloudtrace::v2::TraceService::Stub>
+      trace_service_stub;
 };
 
 class StackdriverExporter {
  public:
   // Registers the exporter.
-  static void Register(const StackdriverOptions& opts);
+  static void Register(StackdriverOptions&& opts);
+
+  // Registers the exporter. Takes ownership of opts.trace_service_stub
+  // and resets it to nullptr.
+  ABSL_DEPRECATED(
+      "Register() without rvalue StackdriverOptions is deprecated and "
+      "will be removed on or after 2020-01-18")
+  static void Register(StackdriverOptions& opts);
 
   // TODO: Retire this:
   ABSL_DEPRECATED(