redis geo server

contest.go

package main

import (
	"bytes"
	"flag"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/garyburd/redigo/redis"
)

const (
	GEO_HASH_NAME = "contest_geo"
	GEO_DICT_NAME = "contest_dict"
	GEO_DIST      = 10000
	GEO_UNIT      = "km"
)

func newPool(addr string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial:        func() (redis.Conn, error) { return redis.Dial("tcp", addr) },
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
	}
}

var (
	pool         *redis.Pool
	redisServer  = flag.String("redisServer", ":6379", "redis server address")
	listenServer = flag.String("listenServer", ":8000", "server listening address")
)

func main() {
	flag.Parse()
	pool = newPool(*redisServer)

	removeLabel := func(s string) string {
		if strings.HasSuffix(s, "N") || strings.HasSuffix(s, "E") {
			return strings.Trim(s, "NE")
		}
		return "-" + strings.Trim(s, "SW")
	}

	http.HandleFunc("/update", func(w http.ResponseWriter, r *http.Request) {
		s := new(bytes.Buffer)
		_, err := s.ReadFrom(r.Body)
		if err != nil {
			log.Println("read from request error ", r.RemoteAddr, err)
			w.WriteHeader(400)
			w.Write([]byte("error when reading from request"))
			return
		}
		ss := strings.Split(s.String(), ",")
		number, n, e := ss[0], removeLabel(ss[2]), removeLabel(ss[3])
		data := strings.Join(ss[4:], ",")
		c := pool.Get()
		defer c.Close()
		_, err = c.Do("GEOADD", GEO_HASH_NAME, n, e, number)
		if err != nil {
			log.Println("GEOADD error ", r.RemoteAddr, number, n, e, err)
			w.WriteHeader(400)
			w.Write([]byte("GEOADD error"))
			return
		}
		_, err = c.Do("HSET", GEO_DICT_NAME, number, data)
		if err != nil {
			log.Println("HSET error ", r.RemoteAddr, number, data)
			w.WriteHeader(400)
			w.Write([]byte("HSET error"))
			return
		}
		w.WriteHeader(200)
	})
	http.HandleFunc("/get", func(w http.ResponseWriter, r *http.Request) {
		s := new(bytes.Buffer)
		_, err := s.ReadFrom(r.Body)
		if err != nil {
			log.Println("read from request error ", r.RemoteAddr, err)
			w.WriteHeader(400)
			w.Write([]byte("error when reading from request"))
			return
		}
		ss := strings.Split(s.String(), ",")
		n, e := removeLabel(ss[0]), removeLabel(ss[1])
		c := pool.Get()
		defer c.Close()
		nums, err := redis.Values(c.Do("GEORADIUS", GEO_HASH_NAME, n, e, GEO_DIST, GEO_UNIT))
		if err != nil {
			log.Println("GEORADIUS of ", n, e, " error ", err)
			w.WriteHeader(400)
			w.Write([]byte("GEORADIUS error"))
			return
		}
		var ret []string
		for _, n := range nums {
			v, err := redis.String(c.Do("hget", GEO_DICT_NAME, n))
			if err != nil {
				log.Println("HGET ", n, " error ", err)
				w.WriteHeader(404)
				w.Write([]byte("HGET error"))
				return
			}
			ret = append(ret, v)
		}
		w.WriteHeader(200)
		w.Write([]byte(strings.Join(ret, " ")))
	})

	fs := http.FileServer(http.Dir("./static"))
	http.Handle("/static/", fs)

	log.Fatal(http.ListenAndServe(*listenServer, nil))
}