Simple Service Discovery In Go with Zookeeper

Implementing Service Registry in Zookeeper

Photo by taichi nakamura on Unsplash

When creating a web service, we may have multiple instance service, the purpose is to make our app can be scaled horizontally. And, the upstream client will ensure to call the downstream services equally. The terms may be called as load-balancing.

There are various tools for service discovery, in this article, I want to tell you how to use Zookeeper, leveraging it’s concept of Znode and Ephemeral Node, to create our own simple round-robin load-balancer.

Key Concepts:

1. There will be 2 Go project: Service (go-zookeeper-registry-service) and Client (go-zookeeper-registry-client)
2. The service app will host an HTTP endpoint (“/hello”) and the client will hit it
3. Inside the service app, it will create an ephemeral node under the pre-defined znode (go-zookeeper-registry-service) with its unique name assigned for each running app and will be assigned with metadata consisting of Host and Port assigned to the app
4. In the client app, it will list down the children (ephemeral nodes) of go-zookeeper-registry-service, and will select (in round-robin based) of each node
5. Client will then retrieve it’s metadata and will connect to the <host>:<port> from the selected child metadata, this will act as the load-balancing mechanism

Pre-requisites

1. Zookeeper Server running

Codes

package main

import (
 "encoding/json"
 "fmt"
 "log"
 "net/http"
 "os"
 "os/signal"
 "strconv"
 "syscall"
 "time"

 "github.com/samuel/go-zookeeper/zk"
)

func main() {
 if len(os.Args) < 2 {
  log.Fatalf("Port number is required as the first argument")
 }

 port, err := strconv.Atoi(os.Args[1])
 if err != nil {
  log.Fatalf("Invalid port number: %v", err)
 }

 conn := connectToZookeeper()
 defer conn.Close()

 parentPath := "/go-zookeeper-registry-service"
 createParentZnodeIfNotExists(conn, parentPath)

 instancePath := createEphemeralZnode(conn, parentPath, port)
 serviceName := instancePath[len(parentPath)+1:]

 go watchZnode(conn, parentPath)
 go startHTTPServer(port, serviceName)

 waitForShutdown()
}

func connectToZookeeper() *zk.Conn {
 conn, _, err := zk.Connect([]string{"localhost:2181"}, time.Second)
 if err != nil {
  log.Fatalf("Unable to connect to Zookeeper: %v", err)
 }
 return conn
}

func createParentZnodeIfNotExists(conn *zk.Conn, parentPath string) {
 exists, _, err := conn.Exists(parentPath)
 if err != nil {
  log.Fatalf("Unable to check if znode exists: %v", err)
 }
 if !exists {
  _, err = conn.Create(parentPath, []byte{}, 0, zk.WorldACL(zk.PermAll))
  if err != nil {
   log.Fatalf("Unable to create parent znode: %v", err)
  }
  fmt.Printf("Created parent znode: %s\n", parentPath)
 }
}

func createEphemeralZnode(conn *zk.Conn, parentPath string, port int) string {
 timestamp := time.Now().Unix()
 instancePath := fmt.Sprintf("%s/%d", parentPath, timestamp)

 data := map[string]interface{}{
  "port": port,
  "host": "127.0.0.1",
 }
 dataBytes, err := json.Marshal(data)
 if err != nil {
  log.Fatalf("Unable to marshal data to JSON: %v", err)
 }

 _, err = conn.Create(instancePath, dataBytes, zk.FlagEphemeral, zk.WorldACL(zk.PermAll))
 if err != nil {
  log.Fatalf("Unable to create ephemeral znode: %v", err)
 }
 fmt.Printf("Created ephemeral znode: %s with data: %s\n", instancePath, string(dataBytes))

 return instancePath
}

func watchZnode(conn *zk.Conn, parentPath string) {
 for {
  children, _, ch, err := conn.ChildrenW(parentPath)
  if err != nil {
   log.Printf("Unable to watch znode: %v", err)
   time.Sleep(time.Second)
   continue
  }

  fmt.Printf("Current znodes: %v\n", children)

  event := <-ch
  fmt.Printf("Znode event: %v\n", event)
 }
}

func waitForShutdown() {
 sig := make(chan os.Signal, 1)
 signal.Notify(sig, os.Interrupt, syscall.SIGTERM)

 <-sig
 fmt.Println("Shutting down gracefully...")

 fmt.Println("Service stopped.")
}

func startHTTPServer(port int, serviceName string) {
 http.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
  fmt.Fprintf(w, "Hi from service %s", serviceName)
 })

 addr := fmt.Sprintf(":%d", port)
 fmt.Printf("Starting HTTP server on port %d\n", port)
 if err := http.ListenAndServe(addr, nil); err != nil {
  log.Fatalf("Unable to start HTTP server: %v", err)
 }
}

Run this by using this command

> go run . <port_name>

//e.g : go run . 8001

Run this app as multiple service, differentiate by the port

example:

// terminal 1
> go run . 8001

// terminal 2
> go run . 8002

// terminal 3
> go run . 8003

Now, run the client app, the code as follow:

package main

import (
 "encoding/json"
 "fmt"
 "io/ioutil"
 "log"
 "net/http"
 "time"

 "github.com/samuel/go-zookeeper/zk"
)

const (
 zkServers      = "127.0.0.1:2181"
 znodePath      = "/go-zookeeper-registry-service"
 sleepDuration  = 1 * time.Second
)

type NodeData struct {
 Host string `json:"host"`
 Port int    `json:"port"`
}

func main() {
 conn, _, err := zk.Connect([]string{zkServers}, time.Second)
 if err != nil {
  log.Fatalf("Unable to connect to Zookeeper: %v", err)
 }
 defer conn.Close()

 var currentIndex int

 for {
  children, err := getChildren(conn)
  if err != nil {
   log.Println(err)
   time.Sleep(sleepDuration)
   continue
  }

  if len(children) == 0 {
   log.Println("No nodes found in znode")
   time.Sleep(sleepDuration)
   continue
  }

  nodeData, err := getNodeData(conn, children[currentIndex%len(children)])
  if err != nil {
   log.Println(err)
   time.Sleep(sleepDuration)
   continue
  }

  if err := makeHTTPRequest(nodeData); err != nil {
   log.Println(err)
   time.Sleep(sleepDuration)
   continue
  }

  currentIndex++
  time.Sleep(sleepDuration)
 }
}

func getChildren(conn *zk.Conn) ([]string, error) {
 children, _, err := conn.Children(znodePath)
 if err != nil {
  return nil, fmt.Errorf("unable to get children of znode: %v", err)
 }
 return children, nil
}

func getNodeData(conn *zk.Conn, node string) (*NodeData, error) {
 data, _, err := conn.Get(znodePath + "/" + node)
 if err != nil {
  return nil, fmt.Errorf("unable to get data of node %s: %v", node, err)
 }

 var nodeData NodeData
 if err := json.Unmarshal(data, &nodeData); err != nil {
  return nil, fmt.Errorf("unable to unmarshal node data: %v", err)
 }

 return &nodeData, nil
}

func makeHTTPRequest(nodeData *NodeData) error {
 url := fmt.Sprintf("http://127.0.0.1:%d/hello", nodeData.Port)
 resp, err := http.Get(url)
 if err != nil {
  return fmt.Errorf("HTTP request failed: %v", err)
 }
 defer resp.Body.Close()

 body, err := ioutil.ReadAll(resp.Body)
 if err != nil {
  return fmt.Errorf("unable to read response body: %v", err)
 }

 log.Printf("Response from %s: %s", url, body)
 return nil
}

> go run .

This code will hit the HTTP Service every one second, and see that it will hit different port in round-robin manner