1 files changed, 162 insertions, 39 deletions

diff --git a/pkg/server/server.go b/pkg/server/server.go
index f9ec602..0c3d2af 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -12,43 +12,51 @@ import (
 	"io"
 )
 
-type handler func (r *request)
+type cmd struct {
+	name string
+	f func (r *request)
+}
 
 type node struct {
-       h handler
-       nodes map[string]node
+	c *cmd
+	m map[string]*node
 }
 
-var handlers = map[string]handler{}
+var cmds = newNode()
 
 type Server struct {
 	listen net.Listener
 	since time.Time
 
 	wg sync.WaitGroup
-	m sync.Mutex
 
-	done bool
+	once sync.Once
+	done chan struct{}
+
+	env env
 
 	nextCid uint64
 }
 
 type client struct {
-	s *Server
-
 	id uint64
 
+	s *Server
+
 	conn net.Conn
 
 	nextRid uint64
 }
 
 type request struct {
+	id uint64
+
 	c *client
 
-	id uint64
+	cmd *cmd
 
 	name string
+	argc int
 	args []string
 
 	out *bytes.Buffer
@@ -62,27 +70,121 @@ func (r *request) String() string {
 	return fmt.Sprintf("request(%d)", r.id)
 }
 
-func setHandler(h handler, names ...string) {
-	var path []string
+func (r *request) Print(v ...interface{}) {
+	fmt.Fprint(r.out, v...)
+}
 
-	for _, s := range names {
-		if _, ok := handlers[s]; ok {
-			err := fmt.Sprintf("handler '%s' already registered at '%s'",
-				s, strings.Join(path, " "))
-			panic(err)
+func (r *request) Printf(format string, v ...interface{}) {
+	fmt.Fprintf(r.out, format, v...)
+}
+
+func (r *request) Println(v ...interface{}) {
+	fmt.Fprintln(r.out, v...)
+}
+
+func (r *request) Fatal(v ...interface{}) {
+	panic(fmt.Sprint(v...))
+}
+
+func (r *request) Fatalf(format string, v ...interface{}) {
+	panic(fmt.Sprintf(format, v...))
+}
+
+func (r *request) Fatalln(v ...interface{}) {
+	panic(fmt.Sprintln(v...))
+}
+
+func (r *request) expect(c ...int) {
+	desc := func (n int) string {
+		var sep string
+
+		if n == 1 {
+			sep = "is"
+		} else {
+			sep = "are"
 		}
 
-		path = append(path, s)
+		return fmt.Sprintf("%d %s expected", n, sep)
+	}
+
+	switch len(c) {
+	case 0:
+		if r.argc > 0 {
+			r.Fatal("args are not expected")
+		}
 
-		handlers[s] = h
-		break
+	case 1:
+		if r.argc < c[0] {
+			r.Fatal("not enough args: ", desc(c[0]))
+		}
+
+		if r.argc > c[0] {
+			r.Fatal("too many args: ", desc(c[0]))
+		}
+
+	case 2:
+		if r.argc < c[0] {
+			r.Fatal("not enough args: at least ", desc(c[0]))
+		}
+
+		if r.argc > c[1] {
+			r.Fatal("too many args: no more than ", desc(c[1]))
+		}
 	}
 }
 
+func newNode() *node {
+	return &node{m: map[string]*node{}}
+}
+
+func setHandler(f func (r *request), path ...string) {
+	node := cmds
+
+	for _, name := range path {
+		v := node.m[name]
+		if v == nil {
+			v = newNode()
+			node.m[name] = v
+		}
+
+		node = v
+	}
+
+	if node.c != nil {
+		s := strings.Join(path, " ")
+		log.Panicf("handler already registered at '%s'", s)
+	}
+
+	node.c = &cmd{
+		name: strings.Join(path, " "),
+		f: f,
+	}
+}
+
+func getHandler(path []string) (*cmd, []string) {
+	node := cmds
+
+	for n, name := range path {
+		node = node.m[name]
+		if node == nil {
+			return nil, nil
+		}
+
+		if node.c != nil {
+			return node.c, path[n + 1:]
+		}
+	}
+
+	return nil, nil
+}
+
 func (s *Server) isDone() bool {
-	s.m.Lock()
-	defer s.m.Unlock()
-	return s.done
+	select {
+	case <- s.done:
+		return true
+	default:
+		return false
+	}
 }
 
 func New() (*Server, error) {
@@ -94,6 +196,7 @@ func New() (*Server, error) {
 	s := &Server{
 		listen: listen,
 		since: time.Now(),
+		done: make(chan struct{}),
 	}
 
 	return s, nil
@@ -124,11 +227,10 @@ func (s *Server) Run() {
 }
 
 func (s *Server) Stop() {
-	s.m.Lock()
-	s.done = true
-	s.m.Unlock()
-
-	s.listen.Close()
+	s.once.Do(func () {
+		close(s.done)
+		s.listen.Close()
+	})
 }
 
 func (s *Server) newClient(conn net.Conn) *client {
@@ -143,14 +245,10 @@ func (s *Server) newClient(conn net.Conn) *client {
 	return c
 }
 
-func (c *client) newRequest(msg string) *request {
-	args := strings.Split(msg, " ")
-
+func (c *client) newRequest() *request {
 	r := &request{
 		c:	c,
 		id:	c.nextRid,
-		name:	args[0],
-		args:	args[1:],
 		out:	bytes.NewBuffer(nil),
 	}
 
@@ -173,14 +271,12 @@ func (c *client) handle() {
 			break
 		}
 
-		msg := string(buf[:nr])
-		r := c.newRequest(msg)
+		query := string(buf[:nr])
 
-		if h, ok := handlers[r.name]; ok {
-			log.Println(c, r, "run:", msg)
-			h(r)
-		} else {
-			fmt.Fprint(r.out, "unknown command")
+		r := c.newRequest()
+
+		if r.parse(query) {
+			r.run()
 		}
 
 		if r.out.Len() == 0 {
@@ -195,6 +291,33 @@ func (c *client) handle() {
 	}
 }
 
+func (r *request) parse(query string) bool {
+	c, args := getHandler(strings.Split(query, " "))
+
+	if c == nil {
+		r.Print("unknown command")
+		return false
+	}
+
+	r.args = args
+	r.argc = len(args)
+	r.cmd = c
+
+	return true
+}
+
+func (r *request) run() {
+	log.Printf("%s %s run [%s] '%s'", r.c, r, r.cmd.name, strings.Join(r.args, " "))
+
+	defer func () {
+		if e := recover(); e != nil {
+			r.Print(e)
+		}
+	}()
+
+	r.cmd.f(r)
+}
+
 func (c *client) close() {
 	log.Println(c, "close")