From 087ba30a90225fbcd79adcf025a5f4f39ea079bf Mon Sep 17 00:00:00 2001
From: Xiang Li <xiang.li@coreos.com>
Date: Thu, 26 Mar 2015 21:10:23 -0700
Subject: [PATCH] server: introduce /v2/migration/snapshot endpoint

---
 server/server.go | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/server/server.go b/server/server.go
index c88fa27d6..da9c45290 100644
--- a/server/server.go
+++ b/server/server.go
@@ -3,8 +3,10 @@ package server
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"net/http/pprof"
+	"path"
 	"strings"
 	"time"
 
@@ -135,6 +137,7 @@ func (s *Server) installV2(r *mux.Router) {
 	s.handleFunc(r2, "/v2/stats/leader", s.GetLeaderStatsHandler).Methods("GET", "HEAD")
 	s.handleFunc(r2, "/v2/stats/store", s.GetStoreStatsHandler).Methods("GET", "HEAD")
 	s.handleFunc(r2, "/v2/speedTest", s.SpeedTestHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/migration/snapshot", s.SnapshotHandler).Methods("GET")
 }
 
 func (s *Server) installMod(r *mux.Router) {
@@ -363,6 +366,41 @@ func (s *Server) SpeedTestHandler(w http.ResponseWriter, req *http.Request) erro
 	return nil
 }
 
+// SnapshotHandler forces etcd store to do a snapshot. If the disk parameter is set, the snapshot
+// will be written to disk at data-dir/index-migrate.snap. Or the snapshot will be returned as
+// http body.
+func (s *Server) SnapshotHandler(w http.ResponseWriter, req *http.Request) error {
+	data, err := s.Store().Save()
+	if err != nil {
+		http.Error(w, "failed to create snapshot: "+err.Error(), http.StatusInternalServerError)
+		log.Warn("Failed to create snapshot:" + err.Error())
+		return nil
+	}
+
+	disk := req.FormValue("disk")
+	if disk == "true" {
+		name := fmt.Sprintf("%d-migrate.snap", s.peerServer.RaftServer().CommitIndex())
+		err = ioutil.WriteFile(path.Join(s.peerServer.RaftServer().Path(), name), data, 0600)
+		if err != nil {
+			http.Error(w, "failed to save snapshot: "+err.Error(), http.StatusInternalServerError)
+			log.Warn("server: failed to save snapshot: " + err.Error())
+			return nil
+		}
+		log.Infof("server: saved snapshot file %s successfully", name)
+		return nil
+	}
+	if disk != "" && disk != "false" {
+		http.Error(w, "invalid parameter: disk="+disk, http.StatusBadRequest)
+		return nil
+	}
+
+	_, err = w.Write(data)
+	if err != nil {
+		log.Warnf("server: failed to write snapshot to %s: %v", req.RemoteAddr, err.Error())
+	}
+	return nil
+}
+
 // Retrieves metrics from bucket
 func (s *Server) GetMetricsHandler(w http.ResponseWriter, req *http.Request) error {
 	(*s.metrics).Dump(w)