Enhance server discovery

Author: ibuildthecloud

pkg/config/types.go

@@ -39,11 +39,11 @@ var (
 
 type Config struct {
 	RuntimeConfig
-	KubernetesVersion string            `json:"kubernetesVersion,omitempty"`
-	RancherVersion    string            `json:"rancherVersion,omitempty"`
-	Server            string            `json:"server,omitempty"`
-	Discovery         map[string]string `json:"discovery,omitempty"`
-	Role              string            `json:"role,omitempty"`
+	KubernetesVersion string           `json:"kubernetesVersion,omitempty"`
+	RancherVersion    string           `json:"rancherVersion,omitempty"`
+	Server            string           `json:"server,omitempty"`
+	Discovery         *DiscoveryConfig `json:"discovery,omitempty"`
+	Role              string           `json:"role,omitempty"`
 
 	RancherValues      map[string]interface{} `json:"rancherValues,omitempty"`
 	PreInstructions    []plan.Instruction     `json:"preInstructions,omitempty"`
@@ -57,6 +57,14 @@ type Config struct {
 	Registries            *registries.Registry `json:"registries,omitempty"`
 }
 
+type DiscoveryConfig struct {
+	Params          map[string]string `json:"params,omitempty"`
+	ExpectedServers int               `json:"expectedServers,omitempty"`
+	// ServerCacheDuration will remember discovered servers for this amount of time.  This
+	// helps with some discovery protocols like mDNS that can be unreliable
+	ServerCacheDuration string `json:"serverCacheDuration,omitempty"`
+}
+
 func Load(path string) (result Config, err error) {
 	var (
 		values = map[string]interface{}{}

pkg/discovery/discovery.go

@@ -7,18 +7,20 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"net"
 	"net/http"
 	"sort"
+	"strconv"
 	"sync"
 	"time"
 
 	"github.com/hashicorp/go-discover"
-	"github.com/rancher/dynamiclistener/server"
 	"github.com/rancher/rancherd/pkg/config"
 	"github.com/rancher/wrangler/pkg/data/convert"
 	"github.com/rancher/wrangler/pkg/randomtoken"
 	"github.com/rancher/wrangler/pkg/slice"
 	"github.com/sirupsen/logrus"
+	"k8s.io/client-go/util/cert"
 
 	// Include kubernetes provider
 	_ "github.com/hashicorp/go-discover/provider/k8s"
@@ -38,10 +40,17 @@ var (
 )
 
 func DiscoverServerAndRole(ctx context.Context, cfg *config.Config) error {
-	if len(cfg.Discovery) == 0 {
+	if cfg.Discovery == nil {
+		if cfg.Server == "" && cfg.Role == "server" && cfg.Token == "" {
+			cfg.Role = "cluster-init"
+		}
 		return nil
 	}
 
+	if cfg.Token == "" {
+		return fmt.Errorf("token is required to be set")
+	}
+
 	server, clusterInit, err := discoverServerAndRole(ctx, cfg)
 	if err != nil {
 		return err
@@ -51,6 +60,7 @@ func DiscoverServerAndRole(ctx context.Context, cfg *config.Config) error {
 	} else if server != "" {
 		cfg.Server = server
 	}
+	logrus.Infof("Using role=%s and server=%s", cfg.Role, cfg.Server)
 	return nil
 
 }
@@ -68,13 +78,18 @@ func discoverServerAndRole(ctx context.Context, cfg *config.Config) (string, boo
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	server, err := newJoinServer(ctx, port)
+	server, err := newJoinServer(ctx, cfg.Discovery.ServerCacheDuration, port)
 	if err != nil {
 		return "", false, err
 	}
 
+	count := cfg.Discovery.ExpectedServers
+	if count == 0 {
+		count = 3
+	}
+
 	for {
-		server, clusterInit := server.loop(ctx, cfg.Discovery, port, discovery)
+		server, clusterInit := server.loop(ctx, count, cfg.Discovery.Params, port, discovery)
 		if clusterInit {
 			return "", true, nil
 		}
@@ -90,15 +105,33 @@ func discoverServerAndRole(ctx context.Context, cfg *config.Config) (string, boo
 	}
 }
 
-func (j *joinServer) loop(ctx context.Context, params map[string]string, port int64, discovery *discover.Discover) (string, bool) {
+func (j *joinServer) addresses(params map[string]string, discovery *discover.Discover) ([]string, error) {
 	addrs, err := discovery.Addrs(discover.Config(params).String(), log.Default())
+	if err != nil {
+		return nil, err
+	}
+
+	var ips []string
+	for _, addr := range addrs {
+		host, _, err := net.SplitHostPort(addr)
+		if err == nil {
+			ips = append(ips, host)
+		} else {
+			ips = append(ips, addr)
+		}
+	}
+
+	return ips, nil
+}
+
+func (j *joinServer) loop(ctx context.Context, count int, params map[string]string, port int64, discovery *discover.Discover) (string, bool) {
+	addrs, err := j.addresses(params, discovery)
 	if err != nil {
 		logrus.Errorf("failed to discover peers to: %v", err)
 		return "", false
 	}
 
-	sort.Strings(addrs)
-	j.setPeers(addrs)
+	addrs = j.setPeers(addrs)
 
 	var (
 		allAgree = true
@@ -128,7 +161,7 @@ func (j *joinServer) loop(ctx context.Context, params map[string]string, port in
 
 		rancherID := resp.Header.Get("X-Cattle-Rancherd-Id")
 		if rancherID == "" {
-			return fmt.Sprintf("https://%s:%d", addr, port), false
+			return fmt.Sprintf("https://%s", net.JoinHostPort(addr, strconv.FormatInt(port, 10))), false
 		}
 		if i == 0 {
 			firstID = rancherID
@@ -148,41 +181,126 @@ func (j *joinServer) loop(ctx context.Context, params map[string]string, port in
 		}
 	}
 
-	if allAgree && len(addrs) > 2 && firstID == j.id {
-		return "", true
+	if firstID != j.id {
+		logrus.Infof("Waiting for peer %s from %v to initialize", addrs[0], addrs)
+		return "", false
 	}
 
-	return "", false
+	if len(addrs) != count {
+		logrus.Infof("Expecting %d servers currently have %v", count, addrs)
+		return "", false
+	}
+
+	if !allAgree {
+		logrus.Infof("All peers %v do not agree on the peer list", addrs)
+		return "", false
+	}
+
+	logrus.Infof("Currently the elected leader %s from peers %v", firstID, addrs)
+	return "", true
 }
 
 type joinServer struct {
-	lock  sync.Mutex
-	id    string
-	peers []string
+	lock          sync.Mutex
+	id            string
+	peers         []string
+	peerSeen      map[string]time.Time
+	cacheDuration time.Duration
 }
 
 type pingResponse struct {
 	Peers []string `json:"peers,omitempty"`
 }
 
-func newJoinServer(ctx context.Context, port int64) (*joinServer, error) {
+func newJoinServer(ctx context.Context, cacheDuration string, port int64) (*joinServer, error) {
 	id, err := randomtoken.Generate()
 	if err != nil {
 		return nil, err
 	}
 
+	if cacheDuration == "" {
+		cacheDuration = "5m"
+	}
+
+	duration, err := time.ParseDuration(cacheDuration)
+	if err != nil {
+		return nil, err
+	}
+
 	j := &joinServer{
-		id: id,
+		id:            id,
+		cacheDuration: duration,
+		peerSeen:      map[string]time.Time{},
+	}
+
+	cert, key, err := cert.GenerateSelfSignedCertKey("rancherd-bootstrap", nil, nil)
+	if err != nil {
+		return nil, err
+	}
+	certs, err := tls.X509KeyPair(cert, key)
+	if err != nil {
+		return nil, err
+	}
+	l, err := tls.Listen("tcp", fmt.Sprintf(":%d", port), &tls.Config{
+		Certificates: []tls.Certificate{
+			certs,
+		},
+	})
+	if err != nil {
+		return nil, err
 	}
+	server := &http.Server{
+		BaseContext: func(_ net.Listener) context.Context {
+			return ctx
+		},
+		Handler: j,
+	}
+	go func() {
+		err := server.Serve(l)
+		if err != nil {
+			logrus.Errorf("failed to server bootstrap http server: %v", err)
+		}
+	}()
+	go func() {
+		<-ctx.Done()
+		server.Shutdown(context.Background())
+		l.Close()
+	}()
 
-	return j, server.ListenAndServe(ctx, int(port), 0, j, nil)
+	return j, nil
 }
 
-func (j *joinServer) setPeers(peers []string) {
+func (j *joinServer) setPeers(peers []string) []string {
 	j.lock.Lock()
 	defer j.lock.Unlock()
-	logrus.Infof("current set of peers: %v", peers)
-	j.peers = peers
+
+	// purge
+	now := time.Now()
+	for k, v := range j.peerSeen {
+		if v.Add(j.cacheDuration).Before(now) {
+			logrus.Info("Forgetting peer %s", k)
+			delete(j.peerSeen, k)
+		}
+	}
+
+	// add
+	for _, peer := range peers {
+		if _, ok := j.peerSeen[peer]; !ok {
+			logrus.Info("New peer discovered %s", peer)
+		}
+		j.peerSeen[peer] = now
+	}
+
+	// sort
+	newPeers := make([]string, 0, len(j.peerSeen))
+	for k := range j.peerSeen {
+		newPeers = append(newPeers, k)
+	}
+	sort.Strings(newPeers)
+
+	j.peers = newPeers
+	logrus.Infof("current set of peers: %v", j.peers)
+	return j.peers
 }
 
 func (j *joinServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {