From 65f5bdec02ef1535b566738c3d285fd89071f965 Mon Sep 17 00:00:00 2001
From: easy <g-easy@users.noreply.github.com>
Date: Wed, 5 Sep 2018 23:51:25 +1000
Subject: [PATCH] Add StackdriverExporterOptions to
 StackdriverExporter::Register().

Use this to configure RPC deadlines and whether to block until
the exporter is connected.

Also change Register() to return a grpc::Status.
---
 .../internal/stackdriver_exporter.cc          | 27 ++++++++++++++-----
 .../trace/stackdriver/stackdriver_exporter.h  | 22 +++++++++++++--
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/opencensus/exporters/trace/stackdriver/internal/stackdriver_exporter.cc b/opencensus/exporters/trace/stackdriver/internal/stackdriver_exporter.cc
index 945c4fe2..c5c472a2 100644
--- a/opencensus/exporters/trace/stackdriver/internal/stackdriver_exporter.cc
+++ b/opencensus/exporters/trace/stackdriver/internal/stackdriver_exporter.cc
@@ -246,10 +246,12 @@ void ConvertSpans(
 class Handler : public ::opencensus::trace::exporter::SpanExporter::Handler {
  public:
   Handler(absl::string_view project_id,
-          const std::shared_ptr<grpc::Channel>& channel)
+          const std::shared_ptr<grpc::Channel>& channel,
+          const StackdriverExporterOptions& opts)
       : project_id_(project_id),
-        stub_(::google::devtools::cloudtrace::v2::TraceService::NewStub(
-            channel)) {}
+        stub_(
+            ::google::devtools::cloudtrace::v2::TraceService::NewStub(channel)),
+        opts_(opts) {}
 
   void Export(const std::vector<::opencensus::trace::exporter::SpanData>& spans)
       override;
@@ -257,6 +259,7 @@ class Handler : public ::opencensus::trace::exporter::SpanExporter::Handler {
  private:
   const std::string project_id_;
   std::unique_ptr<google::devtools::cloudtrace::v2::TraceService::Stub> stub_;
+  const StackdriverExporterOptions opts_;
 };
 
 void Handler::Export(
@@ -267,7 +270,7 @@ void Handler::Export(
   ::google::protobuf::Empty response;
   grpc::ClientContext context;
   context.set_deadline(
-      ConvertToTimespec(absl::Now() + absl::Milliseconds(3000)));
+      ConvertToTimespec(absl::Now() + absl::Seconds(opts_.rpc_deadline_secs)));
   grpc::Status status = stub_->BatchWriteSpans(&context, request, &response);
   if (!status.ok()) {
     std::cerr << "BatchWriteSpans failed: "
@@ -277,11 +280,23 @@ void Handler::Export(
 
 }  // namespace
 
-void StackdriverExporter::Register(absl::string_view project_id) {
+grpc::Status StackdriverExporter::Register(
+    absl::string_view project_id, const StackdriverExporterOptions& opts) {
   auto creds = grpc::GoogleDefaultCredentials();
   auto channel = ::grpc::CreateChannel(kGoogleStackdriverTraceAddress, creds);
+  if (opts.block_until_connected) {
+    if (!channel->WaitForStateChange(
+            GRPC_CHANNEL_IDLE,
+            ConvertToTimespec(absl::Now() +
+                              absl::Seconds(opts.rpc_deadline_secs)))) {
+      // channel->close()?
+      return grpc::Status(grpc::StatusCode::DEADLINE_EXCEEDED,
+                          "initial connection timed out");
+    }
+  }
   ::opencensus::trace::exporter::SpanExporter::RegisterHandler(
-      absl::make_unique<Handler>(project_id, channel));
+      absl::make_unique<Handler>(project_id, channel, opts));
+  return grpc::Status::OK;
 }
 
 }  // namespace trace
diff --git a/opencensus/exporters/trace/stackdriver/stackdriver_exporter.h b/opencensus/exporters/trace/stackdriver/stackdriver_exporter.h
index 0f30e197..86477aac 100644
--- a/opencensus/exporters/trace/stackdriver/stackdriver_exporter.h
+++ b/opencensus/exporters/trace/stackdriver/stackdriver_exporter.h
@@ -16,15 +16,33 @@
 #define OPENCENSUS_EXPORTERS_TRACE_STACKDRIVER_STACKDRIVER_EXPORTER_H_
 
 #include "absl/strings/string_view.h"
+#include "opencensus/common/internal/grpc/status.h"
 
 namespace opencensus {
 namespace exporters {
 namespace trace {
 
+struct StackdriverExporterOptions {
+  // If true, Register() blocks until a connection is established and returns a
+  // meaningful Status object.
+  bool block_until_connected = false;
+
+  // Deadline to set on outgoing RPCs.
+  double rpc_deadline_secs = 3.0;
+
+  static StackdriverExporterOptions Default() {
+    return StackdriverExporterOptions();
+  }
+};
+
 class StackdriverExporter {
  public:
-  // Registers the exporter and sets the project ID.
-  static void Register(absl::string_view project_id);
+  // Registers the exporter and sets the project ID. Returns OK and initializes
+  // in the background, unless opts.block_until_connected is true, in which case
+  // it blocks and returns a meaningful Status.
+  static grpc::Status Register(absl::string_view project_id,
+                               const StackdriverExporterOptions& opts =
+                                   StackdriverExporterOptions::Default());
 
  private:
   StackdriverExporter() = delete;