From c8f97878d2806ab0c34eda60c1c51f24a75fa44b Mon Sep 17 00:00:00 2001
From: Matthias Radestock <matthias.radestock@gmail.com>
Date: Fri, 9 Jun 2017 12:25:48 +0100
Subject: [PATCH] re-target app clients when name resolution changes

Fixes #2578.
---
 probe/appclient/app_client.go        | 23 +++++++++++++++++++++++
 probe/appclient/multi_client.go      |  4 +++-
 probe/appclient/multi_client_test.go |  7 +++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/probe/appclient/app_client.go b/probe/appclient/app_client.go
index 9d218857c6..e207daabcb 100644
--- a/probe/appclient/app_client.go
+++ b/probe/appclient/app_client.go
@@ -31,6 +31,8 @@ type AppClient interface {
 	PipeConnection(string, xfer.Pipe)
 	PipeClose(string) error
 	Publish(r io.Reader) error
+	Target() url.URL
+	ReTarget(url.URL)
 	Stop()
 }
 
@@ -84,10 +86,14 @@ func NewAppClient(pc ProbeConfig, hostname string, target url.URL, control xfer.
 }
 
 func (c *appClient) url(path string) string {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 	return c.target.String() + path
 }
 
 func (c *appClient) wsURL(path string) string {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 	output := c.target //copy the url
 	if output.Scheme == "https" {
 		output.Scheme = "wss"
@@ -140,6 +146,23 @@ func (c *appClient) releaseGoroutine() {
 	c.backgroundWait.Done()
 }
 
+func (c *appClient) Target() url.URL {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	return c.target
+}
+
+// Re-target the appClient, publishing to a new URL. Note that control
+// and pipe websocket connections are left untouched since we don't
+// want to disrupt them just because there's some load-balancing going
+// on. They *will* however pick up the new URL when terminating
+// (e.g. due to errors or when the connection drops).
+func (c *appClient) ReTarget(target url.URL) {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	c.target = target
+}
+
 // Stop stops the appClient.
 func (c *appClient) Stop() {
 	c.mtx.Lock()
diff --git a/probe/appclient/multi_client.go b/probe/appclient/multi_client.go
index c38c9d3edb..f741c3129d 100644
--- a/probe/appclient/multi_client.go
+++ b/probe/appclient/multi_client.go
@@ -103,7 +103,9 @@ func (c *multiClient) Set(hostname string, urls []url.URL) {
 	hostIDs := report.MakeIDList()
 	for tuple := range clients {
 		hostIDs = hostIDs.Add(tuple.ID)
-		if _, ok := c.clients[tuple.ID]; !ok {
+		if client, ok := c.clients[tuple.ID]; ok {
+			client.ReTarget(tuple.AppClient.Target())
+		} else {
 			c.clients[tuple.ID] = tuple.AppClient
 			if !c.noControls {
 				tuple.AppClient.ControlConnection()
diff --git a/probe/appclient/multi_client_test.go b/probe/appclient/multi_client_test.go
index 93795083bb..744a9d05ab 100644
--- a/probe/appclient/multi_client_test.go
+++ b/probe/appclient/multi_client_test.go
@@ -26,6 +26,13 @@ func (c *mockClient) ControlConnection() {
 	c.count++
 }
 
+func (c *mockClient) Target() url.URL {
+	return url.URL{}
+}
+
+func (c *mockClient) ReTarget(_ url.URL) {
+}
+
 func (c *mockClient) Stop() {
 	c.stopped++
 }