Merge branch 'master' into resharding

go/proc/proc.go

@@ -0,0 +1,89 @@
+// Copyright 2012, Google Inc. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package proc allows you to configure servers to be
+// restarted with negligible downtime.
+package proc
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/youtube/vitess/go/relog"
+)
+
+const pidURL = "/debug/pid"
+
+// Listen tries to create a listener on the specified tcp port.
+// Before creating the listener, it checks to see if there is another
+// server already using the port. If there is one, it sends a USR1
+// signal requesting the server to shutdown, and then attempts to
+// to create the listener.
+func Listen(port string) (l net.Listener, err error) {
+	killPredecessor(port)
+	return listen(port)
+}
+
+// Wait creates an HTTP handler on pidURL, and serves the current process
+// pid on it. It then creates a signal handler and waits for SIGTERM or
+// SIGUSR1, and returns when the signal is received. A new server that comes
+// up will query this URL. If it receives a valid response, it will send a
+// SIGUSR1 signal and attempt to bind to the port the current server is using.
+func Wait() os.Signal {
+	http.HandleFunc(pidURL, func(r http.ResponseWriter, req *http.Request) {
+		r.Write(strconv.AppendInt(nil, int64(os.Getpid()), 10))
+	})
+
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGTERM, syscall.SIGUSR1)
+	return <-c
+}
+
+func killPredecessor(port string) {
+	resp, err := http.Get(fmt.Sprintf("http://localhost:%s%s", port, pidURL))
+	if err != nil {
+		if !strings.Contains(err.Error(), "connection refused") {
+			relog.Error("unexpected error on port %v: %v, trying to start anyway", port, err)
+		}
+		return
+	}
+	num, err := ioutil.ReadAll(resp.Body)
+	resp.Body.Close()
+	if err != nil {
+		relog.Error("could not read pid: %vd, trying to start anyway", err)
+		return
+	}
+	pid, err := strconv.Atoi(string(num))
+	if err != nil {
+		relog.Error("could not read pid: %vd, trying to start anyway", err)
+		return
+	}
+	err = syscall.Kill(pid, syscall.SIGUSR1)
+	if err != nil {
+		relog.Error("error killing %v: %v, trying to start anyway", pid, err)
+	}
+}
+
+func listen(port string) (l net.Listener, err error) {
+	for i := 0; i < 100; i++ {
+		l, err = net.Listen("tcp", ":"+port)
+		if err != nil {
+			if strings.Contains(err.Error(), "already in use") {
+				time.Sleep(1 * time.Millisecond)
+				continue
+			}
+			return nil, err
+		}
+		break
+	}
+	return l, err
+}

go/proc/proc_test.go

@@ -0,0 +1,121 @@
+// Copyright 2012, Google Inc. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package proc
+
+import (
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"os/exec"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+)
+
+var port = flag.String("port", "", "http port")
+
+func TestRestart(t *testing.T) {
+	switch os.Getenv("SERVER_NUM") {
+	case "":
+		testLaunch(t)
+	case "1":
+		testServer(t, syscall.SIGUSR1)
+	case "2":
+		testServer(t, syscall.SIGTERM)
+	}
+}
+
+var testPort = "12345"
+
+func testLaunch(t *testing.T) {
+	var err error
+
+	cmd1 := launchServer(t, 1)
+	defer cmd1.Process.Kill()
+	testPid(t, cmd1.Process.Pid)
+
+	cmd2 := launchServer(t, 2)
+	defer cmd2.Process.Kill()
+	err = cmd1.Wait()
+	if err != nil {
+		t.Error(err)
+	}
+	testPid(t, cmd2.Process.Pid)
+
+	err = syscall.Kill(cmd2.Process.Pid, syscall.SIGTERM)
+	if err != nil {
+		t.Error(err)
+	}
+	err = cmd2.Wait()
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func launchServer(t *testing.T, num int) *exec.Cmd {
+	cmd := exec.Command(os.Args[0], "-test.run=^TestRestart$", "-port", testPort)
+	cmd.Env = []string{fmt.Sprintf("SERVER_NUM=%d", num)}
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Start()
+	if err != nil {
+		t.Fatal(err)
+	}
+	return cmd
+}
+
+func testPid(t *testing.T, want int) {
+	var resp *http.Response
+	var err error
+	for i := 0; i < 20; i++ {
+		resp, err = http.Get(fmt.Sprintf("http://localhost:%s%s", testPort, pidURL))
+		if err != nil {
+			if i == 19 {
+				t.Fatal(err)
+			}
+			if strings.Contains(err.Error(), "connection refused") {
+				time.Sleep(1000 * time.Millisecond)
+				continue
+			}
+			t.Fatalf("unexpected error on port %v: %v", testPort, err)
+		}
+		break
+	}
+	num, err := ioutil.ReadAll(resp.Body)
+	resp.Body.Close()
+	if err != nil {
+		t.Fatalf("could not read pid: %vd", err)
+	}
+	got, err := strconv.Atoi(string(num))
+	if err != nil {
+		t.Fatalf("could not read pid: %vd", err)
+	}
+	if want != got {
+		t.Errorf("want %d, got %d", want, got)
+	}
+}
+
+func testServer(t *testing.T, want syscall.Signal) {
+	flag.Parse()
+	if *port != "12345" {
+		t.Errorf("want 12345, got %s", *port)
+	}
+
+	l, err := Listen(*port)
+	if err != nil {
+		t.Fatalf("could not initialize listener: %v", err)
+	}
+	go http.Serve(l, nil)
+	got := Wait()
+	l.Close()
+	if want != got {
+		t.Errorf("want %v, got %v", want, got)
+	}
+}