test github.com/weibocom/nginx-upsync-module

auto_reg.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"

	log "github.com/Sirupsen/logrus"
)

var (
	autoReg     bool
	cli         = http.DefaultClient
	bind        = flag.String("bind", ":8080", "web app bind to addr")
	consulAddr  = flag.String("consul", "", "consul listening addr")
	serviceName = flag.String("service", "", "name of upstream/service registered in the consul")
)

type fct func() error

type server struct {
	name       string
	beforeFunc chan fct
	afterFunc  chan func()
}

func newServer(name string) *server {
	s := &server{
		name:       name,
		beforeFunc: make(chan fct, 1),
		afterFunc:  make(chan func()),
	}
	go func() {
		beforeFunc := <-s.beforeFunc
		if err := beforeFunc(); err != nil {
			log.Fatalln(err)
		}
		afterFunc := <-s.afterFunc
		go afterFunc()
	}()
	return s
}

func (s *server) hello(w http.ResponseWriter, req *http.Request) {
	w.Write([]byte(s.name))
}

func doHTTPRequest(method, url, body string) error {
	var buf io.Reader
	if body != "" {
		buf = bytes.NewBuffer([]byte(body))
	}
	if method == "" {
		method = "GET"
	}

	req, err := http.NewRequest(method, url, buf)
	if err != nil {
		return err
	}
	resp, err := cli.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode > 400 {
		return fmt.Errorf("status code: %s", resp.Status)
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	log.Infof("HTTP Request: %s %s %s", method, url, string(b))
	return nil
}

func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, addr := range addrs {
		if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	return ""
}

func main() {
	flag.Parse()

	s := newServer("test github.com/weibocom/nginx-upsync-module")
	http.HandleFunc("/", s.hello)
	if err := http.ListenAndServe(*bind, nil); err != nil {
		log.Fatalln(err)
	}

	autoReg = *consulAddr != "" && *serviceName != ""
	if autoReg {
		rawURL := *consulAddr + "/v1/kv/upstreams/" + *serviceName + "/" + getLocalIP() + *bind
		uri, err := url.Parse(rawURL)
		if err != nil {
			log.Fatalln(err)
		}
		s.beforeFunc <- func() error {
			return doHTTPRequest("PUT", uri.String(), "")
		}
		s.afterFunc <- func() {
			ch := make(chan os.Signal, 1)
			signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
			select {
			case s := <-ch:
				if err := doHTTPRequest("DELETE", uri.String(), ""); err != nil {
					log.Fatalln(err)
				}
				log.Infof("Signal %s, QUIT...", s.String())
				os.Exit(1)
			}
		}
	}
}