Merge pull request #5 from myENA/feature/addr-update

Feature/addr update

Author: dcarbone

README.md

@@ -23,6 +23,9 @@ Option      | Description
 `-f`        | Force killing of all matches, including healthy services
 `-s string` | Limit search by service address (regexp)
 `-t string` | Limit search by tag
+`-local-addr string` | Address with port of "local" agent to use to retrieve service list from (defaults to value of `CONSUL_HTTP_ADDR` env var)
+`-remote-port int` | Port to use when connecting to remote agents.  Defaults to `8500`
+`-token string` | ACL token to use in all API requests
 `-v`        | Verbose
 `-vv`       | Increased verbosity
 `-vvv`      | Super verbosity

zombie.go main.gozombie.go renamed to main.go

@@ -6,13 +6,16 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"regexp"
 
 	"github.com/hashicorp/consul/api"
 	"github.com/olekukonko/tablewriter"
 )
 
 // this is the default port for talking to remote consul agents
-const defaultPort = 8500
+const (
+	defaultPort = 8500
+)
 
 type verbosityLevel uint8
 
@@ -38,15 +41,16 @@ func usage(code int) {
 }
 
 func main() {
-	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
 	serviceString := fs.String("s", "", "Limit search by service address (regexp)")
 	tag := fs.String("t", "", "Limit search by tag")
 	force := fs.Bool("f", false, "Force killing of all matches, including healthy services")
-	port := fs.Int("port", defaultPort, "Port to use when connecting to remote agents")
-	token := fs.String("token", "", "Token to use when connecting to remote agents")
+	localAddr := fs.String("local-addr", os.Getenv(api.HTTPAddrEnvName), "Address with port of \"local\" agent.  Used to list services.")
+	remotePort := fs.Int("remote-port", defaultPort, "Port to use when connecting to remote agents")
+	token := fs.String("token", os.Getenv(api.HTTPTokenEnvName), "ACL token.  Used in all api queries.")
 	v1 := fs.Bool("v", false, "Verbose")
-	v2 := fs.Bool("vv", false, "Increased Verbosity")
-	v3 := fs.Bool("vvv", false, "Super Verbosity")
+	v2 := fs.Bool("vv", false, "Really verbose")
+	v3 := fs.Bool("vvv", false, "Extremely verbose")
 	if err := fs.Parse(os.Args[1:]); err != nil {
 		log.Printf("Error parsing args: %s", err)
 		os.Exit(1)
@@ -71,19 +75,111 @@ func main() {
 	switch cmd {
 	// define a couple synonyms to "hunt" as well
 	case "hunt", "find", "search":
-		serviceList := getList(*serviceString, *token, *tag)
+		serviceList := getList(*localAddr, *token, *serviceString, *tag)
 		printList(serviceList, verbosity)
 
 	case "kill":
-		serviceList := getList(*serviceString, *token, *tag)
-		deregister(serviceList, *port, *token, *force)
+		serviceList := getList(*localAddr, *token, *serviceString, *tag)
+		deregister(*remotePort, *token, serviceList, *force)
 
 	default:
 		usage(1)
 	}
 
 }
 
+// get a client handle for a specified address (or the local agent if "")
+func getClient(address, token string) (*api.Client, error) {
+	config := api.DefaultConfig()
+	if address != "" {
+		config.Address = address
+	}
+	if token != "" {
+		config.Token = token
+	}
+	return api.NewClient(config)
+}
+
+// get a list of all services, limit to those matching the search criteria
+func getList(localAddr, token, serviceString, tag string) []*api.ServiceEntry {
+	client, err := getClient(localAddr, token)
+	if err != nil {
+		log.Fatalf("Unable to get a consul client connection: %s\n", err)
+	}
+
+	serviceList, _, err := client.Catalog().Services(nil)
+	if err != nil {
+		log.Fatalf("Unable to get list of services from catalog: %s", err)
+	}
+
+	nodeServiceMap := make(map[string]map[string]*api.ServiceEntry)
+
+	for svc := range serviceList {
+		entries, _, err := client.Health().Service(svc, tag, false, nil)
+		if err != nil {
+			log.Fatalf("Unable to query for service \"%s\" health: %s", svc, err)
+		}
+		for _, entry := range entries {
+			if _, ok := nodeServiceMap[entry.Node.Node]; !ok {
+				nodeServiceMap[entry.Node.Node] = make(map[string]*api.ServiceEntry)
+			}
+			nodeServiceMap[entry.Node.Node][entry.Service.ID] = entry
+		}
+	}
+
+	seOut := make([]*api.ServiceEntry, 0)
+	if serviceString == "" {
+		for _, services := range nodeServiceMap {
+			for _, service := range services {
+				seOut = append(seOut, service)
+			}
+		}
+	} else {
+		var (
+			serviceID string
+			services  map[string]*api.ServiceEntry
+			service   *api.ServiceEntry
+		)
+		re := regexp.MustCompile(serviceString)
+		for _, services = range nodeServiceMap {
+			for serviceID, service = range services {
+				if re.FindString(serviceID) != "" || re.FindString(service.Service.Service) != "" {
+					seOut = append(seOut, service)
+				}
+			}
+		}
+	}
+
+	return seOut
+}
+
+// A service entry is considered healty if all the eligible checks are passing.
+// serfChecks are not eligible
+func isHealthy(se *api.ServiceEntry) bool {
+	if se == nil || se.Checks == nil || len(se.Checks) == 0 {
+		// No checks = failing
+		return false
+	}
+
+	healthy := true
+	eligible := 0
+	for _, c := range se.Checks {
+		if c.Name == "serfHealth" {
+			continue
+		}
+		// All found checks have to be passing
+		healthy = healthy && (c.Status == api.HealthPassing)
+		eligible++
+	}
+
+	// No eligible checks were found
+	if eligible == 0 {
+		return false
+	}
+
+	return healthy
+}
+
 // display a list of matching services
 func printList(serviceList []*api.ServiceEntry, v verbosityLevel) {
 	var header, footer []string
@@ -142,10 +238,10 @@ func printList(serviceList []*api.ServiceEntry, v verbosityLevel) {
 }
 
 // kill those services that are failing in the passed list, or all if force is true
-func deregister(serviceList []*api.ServiceEntry, port int, token string, force bool) {
+func deregister(remotePort int, token string, serviceList []*api.ServiceEntry, force bool) {
 	for _, se := range serviceList {
 		if !isHealthy(se) || force {
-			fullAddress := fmt.Sprintf("%s:%d", se.Node.Address, port)
+			fullAddress := fmt.Sprintf("%s:%d", se.Node.Address, remotePort)
 			fmt.Printf("Deregistering %s: %s (%s)\n", se.Service.Service, se.Service.ID, fullAddress)
 			client, err := getClient(fullAddress, token)
 			if err != nil {