From 9ea09cc41ca90ba57b9ba304ab5e5e879eff82fd Mon Sep 17 00:00:00 2001
From: Konstantin Lebedev <9497591+kmlebedev@users.noreply.github.com>
Date: Wed, 16 Feb 2022 14:18:36 +0500
Subject: [PATCH] healthz check to avoid drain pod with last replicas

---
 weed/server/volume_server.go                |  1 +
 weed/server/volume_server_handlers_admin.go | 19 +++++++++++++++++++
 weed/topology/store_replicate.go            |  6 +++---
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/weed/server/volume_server.go b/weed/server/volume_server.go
index 2551cc6e6..4199ae36b 100644
--- a/weed/server/volume_server.go
+++ b/weed/server/volume_server.go
@@ -98,6 +98,7 @@ func NewVolumeServer(adminMux, publicMux *http.ServeMux, ip string,
 
 	handleStaticResources(adminMux)
 	adminMux.HandleFunc("/status", vs.statusHandler)
+	adminMux.HandleFunc("/healthz", vs.healthzHandler)
 	if signingKey == "" || enableUiAccess {
 		// only expose the volume server details for safe environments
 		adminMux.HandleFunc("/ui/index.html", vs.uiStatusHandler)
diff --git a/weed/server/volume_server_handlers_admin.go b/weed/server/volume_server_handlers_admin.go
index 7e6c06871..37cf109e2 100644
--- a/weed/server/volume_server_handlers_admin.go
+++ b/weed/server/volume_server_handlers_admin.go
@@ -1,6 +1,7 @@
 package weed_server
 
 import (
+	"github.com/chrislusf/seaweedfs/weed/topology"
 	"net/http"
 	"path/filepath"
 
@@ -9,6 +10,24 @@ import (
 	"github.com/chrislusf/seaweedfs/weed/util"
 )
 
+func (vs *VolumeServer) healthzHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Server", "SeaweedFS Volume "+util.VERSION)
+	volumeInfos := vs.store.VolumeInfos()
+	for _, vinfo := range volumeInfos {
+		if len(vinfo.Collection) == 0 {
+			continue
+		}
+		if vinfo.ReplicaPlacement.GetCopyCount() > 1 {
+			_, err := topology.GetWritableRemoteReplications(vs.store, vs.grpcDialOption, vinfo.Id, vs.GetMaster)
+			if err != nil {
+				w.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+		}
+	}
+	w.WriteHeader(http.StatusOK)
+}
+
 func (vs *VolumeServer) statusHandler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Server", "SeaweedFS Volume "+util.VERSION)
 	m := make(map[string]interface{})
diff --git a/weed/topology/store_replicate.go b/weed/topology/store_replicate.go
index 65625b3b7..7bb10f1da 100644
--- a/weed/topology/store_replicate.go
+++ b/weed/topology/store_replicate.go
@@ -29,7 +29,7 @@ func ReplicatedWrite(masterFn operation.GetMasterFn, grpcDialOption grpc.DialOpt
 	var remoteLocations []operation.Location
 	if r.FormValue("type") != "replicate" {
 		// this is the initial request
-		remoteLocations, err = getWritableRemoteReplications(s, grpcDialOption, volumeId, masterFn)
+		remoteLocations, err = GetWritableRemoteReplications(s, grpcDialOption, volumeId, masterFn)
 		if err != nil {
 			glog.V(0).Infoln(err)
 			return
@@ -114,7 +114,7 @@ func ReplicatedDelete(masterFn operation.GetMasterFn, grpcDialOption grpc.DialOp
 
 	var remoteLocations []operation.Location
 	if r.FormValue("type") != "replicate" {
-		remoteLocations, err = getWritableRemoteReplications(store, grpcDialOption, volumeId, masterFn)
+		remoteLocations, err = GetWritableRemoteReplications(store, grpcDialOption, volumeId, masterFn)
 		if err != nil {
 			glog.V(0).Infoln(err)
 			return
@@ -174,7 +174,7 @@ func DistributedOperation(locations []operation.Location, op func(location opera
 	return ret.Error()
 }
 
-func getWritableRemoteReplications(s *storage.Store, grpcDialOption grpc.DialOption, volumeId needle.VolumeId, masterFn operation.GetMasterFn) (remoteLocations []operation.Location, err error) {
+func GetWritableRemoteReplications(s *storage.Store, grpcDialOption grpc.DialOption, volumeId needle.VolumeId, masterFn operation.GetMasterFn) (remoteLocations []operation.Location, err error) {
 
 	v := s.GetVolume(volumeId)
 	if v != nil && v.ReplicaPlacement.GetCopyCount() == 1 {