From 9f70568a02bc63d53f558789dca7176db802d0ac Mon Sep 17 00:00:00 2001
From: Yicheng Qin <qycqycqycqycqyc@gmail.com>
Date: Wed, 21 Jan 2015 21:46:47 -0800
Subject: [PATCH] etcd: register usable versions when bootstrap

---
 etcd/etcd.go    |  1 +
 etcd/upgrade.go | 59 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 etcd/upgrade.go

diff --git a/etcd/etcd.go b/etcd/etcd.go
index b23fa3db5..20c9122c5 100644
--- a/etcd/etcd.go
+++ b/etcd/etcd.go
@@ -309,6 +309,7 @@ func (e *Etcd) runServer() {
 	for {
 		if e.mode == PeerMode {
 			log.Infof("%v starting in peer mode", e.Config.Name)
+			go registerAvailableInternalVersions(e.Config.Name, e.Config.Addr, e.Config.EtcdTLSInfo())
 			// Starting peer server should be followed close by listening on its port
 			// If not, it may leave many requests unaccepted, or cannot receive heartbeat from the cluster.
 			// One severe problem caused if failing receiving heartbeats is when the second node joins one-node cluster,
diff --git a/etcd/upgrade.go b/etcd/upgrade.go
new file mode 100644
index 000000000..8cb2319be
--- /dev/null
+++ b/etcd/upgrade.go
@@ -0,0 +1,59 @@
+package etcd
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"time"
+
+	"github.com/coreos/etcd/log"
+	"github.com/coreos/etcd/server"
+	"github.com/coreos/etcd/third_party/github.com/coreos/go-etcd/etcd"
+)
+
+var defaultEtcdBinaryDir = "/usr/libexec/etcd/internal_versions/"
+
+func registerAvailableInternalVersions(name string, addr string, tls *server.TLSInfo) {
+	var c *etcd.Client
+	if tls.Scheme() == "http" {
+		c = etcd.NewClient([]string{addr})
+	} else {
+		var err error
+		c, err = etcd.NewTLSClient([]string{addr}, tls.CertFile, tls.KeyFile, tls.CAFile)
+		if err != nil {
+			log.Fatalf("client TLS error: %v", err)
+		}
+	}
+
+	vers, err := getInternalVersions()
+	if err != nil {
+		log.Infof("failed to get local etcd versions: %v", err)
+		return
+	}
+	for _, v := range vers {
+		for {
+			_, err := c.Set("/_etcd/available-internal-versions/"+v+"/"+name, "ok", 0)
+			if err == nil {
+				break
+			}
+			time.Sleep(time.Second)
+		}
+	}
+	log.Infof("%s: available_internal_versions %s is registered into key space successfully.", name, vers)
+}
+
+func getInternalVersions() ([]string, error) {
+	if runtime.GOOS != "linux" {
+		return nil, fmt.Errorf("unmatched os version %v", runtime.GOOS)
+	}
+	etcdBinaryDir := os.Getenv("ETCD_BINARY_DIR")
+	if etcdBinaryDir == "" {
+		etcdBinaryDir = defaultEtcdBinaryDir
+	}
+	dir, err := os.Open(etcdBinaryDir)
+	if err != nil {
+		return nil, err
+	}
+	defer dir.Close()
+	return dir.Readdirnames(-1)
+}