From 8c72a2671bdb468ee86092a3c4230b3771eef125 Mon Sep 17 00:00:00 2001
From: Dmitry Vorobev <dimahabr@gmail.com>
Date: Fri, 28 Apr 2017 18:20:27 +0300
Subject: [PATCH] cli: Add --listening-url-file flag

---
 pkg/base/test_server_args.go |  1 +
 pkg/cli/cliflags/flags.go    |  7 +++++++
 pkg/cli/flags.go             |  2 ++
 pkg/server/config.go         |  4 ++++
 pkg/server/server.go         | 14 +++++++++++++
 pkg/server/server_test.go    | 40 ++++++++++++++++++++++++++++++++++++
 pkg/server/testserver.go     |  4 ++++
 7 files changed, 72 insertions(+)

diff --git a/pkg/base/test_server_args.go b/pkg/base/test_server_args.go
index 82ce70eb9324..6ed0a92e2047 100644
--- a/pkg/base/test_server_args.go
+++ b/pkg/base/test_server_args.go
@@ -67,6 +67,7 @@ type TestServerArgs struct {
 	TimeSeriesQueryWorkerMax int
 	SQLMemoryPoolSize        int64
 	SendNextTimeout          time.Duration
+	ListeningURLFile         string
 
 	// If set, this will be appended to the Postgres URL by functions that
 	// automatically open a connection to the server. That's equivalent to running
diff --git a/pkg/cli/cliflags/flags.go b/pkg/cli/cliflags/flags.go
index aaeccd7a22aa..fbecccd5e2e2 100644
--- a/pkg/cli/cliflags/flags.go
+++ b/pkg/cli/cliflags/flags.go
@@ -221,6 +221,13 @@ communication; it must resolve from other nodes in the cluster.`,
 		Description: `The port to bind to for HTTP requests.`,
 	}
 
+	ListeningURLFile = FlagInfo{
+		Name: "listening-url-file",
+		Description: `
+After the CockroachDB node has started up successfully, it will
+write its connection URL to the specified file.`,
+	}
+
 	PIDFile = FlagInfo{
 		Name: "pid-file",
 		Description: `
diff --git a/pkg/cli/flags.go b/pkg/cli/flags.go
index 9315426ce717..1b50dc2bda31 100644
--- a/pkg/cli/flags.go
+++ b/pkg/cli/flags.go
@@ -247,6 +247,8 @@ func init() {
 		stringFlag(f, &serverCfg.SocketFile, cliflags.Socket, "")
 		_ = f.MarkHidden(cliflags.Socket.Name)
 
+		stringFlag(f, &serverCfg.ListeningURLFile, cliflags.ListeningURLFile, "")
+
 		stringFlag(f, &serverCfg.PIDFile, cliflags.PIDFile, "")
 
 		// Use a separate variable to store the value of ServerInsecure.
diff --git a/pkg/server/config.go b/pkg/server/config.go
index ae509a1838a9..e7792cf8ba43 100644
--- a/pkg/server/config.go
+++ b/pkg/server/config.go
@@ -183,6 +183,10 @@ type Config struct {
 	// actions.
 	EventLogEnabled bool
 
+	// ListeningURLFile indicates the file to which the server writes
+	// its listening URL when it is ready.
+	ListeningURLFile string
+
 	// PIDFile indicates the file to which the server writes its PID when
 	// it is ready.
 	PIDFile string
diff --git a/pkg/server/server.go b/pkg/server/server.go
index ac5e4bbf0089..b60db168461f 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -26,6 +26,7 @@ import (
 	"math"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"strings"
 	"sync"
@@ -48,6 +49,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/migrations"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/server/status"
 	"github.com/cockroachdb/cockroach/pkg/sql"
@@ -858,6 +860,18 @@ func (s *Server) Start(ctx context.Context) error {
 			log.Error(ctx, err)
 		}
 	}
+
+	if s.cfg.ListeningURLFile != "" {
+		pgURL, err := s.cfg.PGURL(url.User(security.RootUser))
+		if err == nil {
+			err = ioutil.WriteFile(s.cfg.ListeningURLFile, []byte(fmt.Sprintf("%s\n", pgURL)), 0644)
+		}
+
+		if err != nil {
+			log.Error(ctx, err)
+		}
+	}
+
 	if err := sdnotify.Ready(); err != nil {
 		log.Errorf(ctx, "failed to signal readiness using systemd protocol: %s", err)
 	}
diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go
index 80434aa059e4..7f0a3c1cd91d 100644
--- a/pkg/server/server_test.go
+++ b/pkg/server/server_test.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"io"
+	"io/ioutil"
 	"net"
 	"net/http"
 	"net/url"
@@ -478,3 +479,42 @@ func TestClusterStores(t *testing.T) {
 		return nil
 	})
 }
+
+func TestListenURLFileCreation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	file, err := ioutil.TempFile(os.TempDir(), t.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := file.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	s, err := serverutils.StartServerRaw(base.TestServerArgs{
+		ListeningURLFile: file.Name(),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Stopper().Stop(context.TODO())
+	defer func() {
+		if err := os.Remove(file.Name()); err != nil {
+			t.Error(err)
+		}
+	}()
+
+	data, err := ioutil.ReadFile(file.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	u, err := url.Parse(string(data))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if s.ServingAddr() != u.Host {
+		t.Fatalf("expected URL %s to match host %s", u, s.ServingAddr())
+	}
+}
diff --git a/pkg/server/testserver.go b/pkg/server/testserver.go
index c1e608ea1dc7..d6a9784cf7a4 100644
--- a/pkg/server/testserver.go
+++ b/pkg/server/testserver.go
@@ -151,6 +151,10 @@ func makeTestConfigFromParams(params base.TestServerArgs) Config {
 		cfg.HTTPAddr = params.HTTPAddr
 	}
 
+	if params.ListeningURLFile != "" {
+		cfg.ListeningURLFile = params.ListeningURLFile
+	}
+
 	// Ensure we have the correct number of engines. Add in-memory ones where
 	// needed. There must be at least one store/engine.
 	if len(params.StoreSpecs) == 0 {