From c1637b49784d5da9da98c97d6dd69b2116e92260 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Fri, 7 Feb 2014 12:11:34 -0800
Subject: [PATCH] agent: Adding server up callback to make state sync faster

---
 command/agent/agent.go | 10 ++++++++--
 command/agent/local.go |  9 +++++++--
 consul/client.go       |  5 +++++
 consul/config.go       |  4 ++++
 consul/serf.go         |  5 +++++
 5 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index ba0b24cd29..28def6f287 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -93,14 +93,17 @@ func Create(config *Config, logOutput io.Writer) (*Agent, error) {
 		shutdownCh:    make(chan struct{}),
 	}
 
+	// Initialize the local state
+	agent.state.Init(config, agent.logger)
+
 	// Setup either the client or the server
 	var err error
 	if config.Server {
 		err = agent.setupServer()
-		agent.state.Init(config, agent.server, agent.logger)
+		agent.state.SetIface(agent.server)
 	} else {
 		err = agent.setupClient()
-		agent.state.Init(config, agent.client, agent.logger)
+		agent.state.SetIface(agent.client)
 	}
 	if err != nil {
 		return nil, err
@@ -162,6 +165,9 @@ func (a *Agent) consulConfig() *consul.Config {
 		base.Bootstrap = true
 	}
 
+	// Setup the ServerUp callback
+	base.ServerUp = a.state.ConsulServerUp
+
 	// Setup the loggers
 	base.LogOutput = a.logOutput
 	return base
diff --git a/command/agent/local.go b/command/agent/local.go
index fbda1ab8eb..6b5854bfbe 100644
--- a/command/agent/local.go
+++ b/command/agent/local.go
@@ -56,9 +56,8 @@ type localState struct {
 }
 
 // Init is used to initialize the local state
-func (l *localState) Init(config *Config, iface consul.Interface, logger *log.Logger) {
+func (l *localState) Init(config *Config, logger *log.Logger) {
 	l.config = config
-	l.iface = iface
 	l.logger = logger
 	l.services = make(map[string]*structs.NodeService)
 	l.serviceStatus = make(map[string]syncStatus)
@@ -68,6 +67,12 @@ func (l *localState) Init(config *Config, iface consul.Interface, logger *log.Lo
 	l.triggerCh = make(chan struct{}, 1)
 }
 
+// SetIface is used to set the Consul interface. Must be set prior to
+// starting anti-entropy
+func (l *localState) SetIface(iface consul.Interface) {
+	l.iface = iface
+}
+
 // changeMade is used to trigger an anti-entropy run
 func (l *localState) changeMade() {
 	select {
diff --git a/consul/client.go b/consul/client.go
index 612d23c376..a00bc9a6e2 100644
--- a/consul/client.go
+++ b/consul/client.go
@@ -218,6 +218,11 @@ func (c *Client) nodeJoin(me serf.MemberEvent) {
 			c.consuls = append(c.consuls, addr)
 		}
 		c.consulLock.Unlock()
+
+		// Trigger the callback
+		if c.config.ServerUp != nil {
+			c.config.ServerUp()
+		}
 	}
 }
 
diff --git a/consul/config.go b/consul/config.go
index 4f0d3cc37c..35de3789ce 100644
--- a/consul/config.go
+++ b/consul/config.go
@@ -64,6 +64,10 @@ type Config struct {
 	// LogOutput is the location to write logs to. If this is not set,
 	// logs will go to stderr.
 	LogOutput io.Writer
+
+	// ServerUp callback can be used to trigger a notification that
+	// a Consul server is now up and known about.
+	ServerUp func()
 }
 
 // DefaultConfig is used to return a sane default configuration
diff --git a/consul/serf.go b/consul/serf.go
index 3dff0f9e2e..61e92b4b37 100644
--- a/consul/serf.go
+++ b/consul/serf.go
@@ -95,6 +95,11 @@ func (s *Server) remoteJoin(me serf.MemberEvent) {
 			s.remoteConsuls[parts.Datacenter] = append(existing, addr)
 		}
 		s.remoteLock.Unlock()
+
+		// Trigger the callback
+		if s.config.ServerUp != nil {
+			s.config.ServerUp()
+		}
 	}
 }