server.go

package redeo

import (
	"context"
	"io"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/mason-leap-lab/redeo/resp"
)

const (
	ASYNC_CB_NAME = "callback"
)

// Server configuration
type Server struct {
	config *Config
	info   *ServerInfo

	cmds map[string]interface{}
	mu   sync.RWMutex
	//	released *sync.WaitGroup
}

// NewServer creates a new server instance
func NewServer(config *Config) *Server {
	if config == nil {
		config = new(Config)
	}

	return &Server{
		config: config,
		info:   newServerInfo(),
		cmds:   make(map[string]interface{}),
	}
}

// Info returns the server info registry
func (srv *Server) Info() *ServerInfo { return srv.info }

// Handle registers a handler for a command.
func (srv *Server) Handle(name string, h Handler) {
	srv.mu.Lock()
	srv.cmds[strings.ToLower(name)] = h
	srv.mu.Unlock()
}

// HandleFunc registers a handler func for a command.
func (srv *Server) HandleFunc(name string, fn HandlerFunc) {
	srv.Handle(name, fn)
}

// HandleStream registers a handler for a streaming command.
func (srv *Server) HandleStream(name string, h StreamHandler) {
	srv.mu.Lock()
	srv.cmds[strings.ToLower(name)] = h
	srv.mu.Unlock()
}

// HandleStreamFunc registers a handler func for a command
func (srv *Server) HandleStreamFunc(name string, fn StreamHandlerFunc) {
	srv.HandleStream(name, fn)
}

// Callback registers a handler for a callback command.
func (srv *Server) HandleCallback(cb Callback) {
	srv.mu.Lock()
	srv.cmds[ASYNC_CB_NAME] = cb
	srv.mu.Unlock()
}

// CallbackFunc registers a handler func for a command
func (srv *Server) HandleCallbackFunc(cbf CallbackFunc) {
	srv.HandleCallback(cbf)
}

// Serve accepts incoming connections on a listener, creating a
// new service goroutine for each.
func (srv *Server) Serve(lis net.Listener) error {
	return srv.serveImpl(lis, true)
}

func (srv *Server) ServeAsync(lis net.Listener) error {
	return srv.serveImpl(lis, false)
}

func (srv *Server) serveImpl(lis net.Listener, sync bool) error {
	for {
		cn, err := lis.Accept()
		if err != nil {
			return err
		}

		if ka := srv.config.TCPKeepAlive; ka > 0 {
			if tc, ok := cn.(*net.TCPConn); ok {
				tc.SetKeepAlive(true)
				tc.SetKeepAlivePeriod(ka)
			}
		}

		go srv.serveClient(newClient(cn), sync)
	}
}

func (srv *Server) GetClient(id uint64) (*Client, bool) {
	return srv.info.Client(id)
}

func (srv *Server) Close(lis net.Listener) {
	lis.Close()
}

func (srv *Server) Release() {
	// srv.mu.Lock()
	// srv.released = &sync.WaitGroup{}
	// srv.mu.Unlock()

	for _, client := range srv.info.Clients() {
		// srv.released.Add(1)
		client.Close()
	}
	// srv.released.Wait()
	// srv.released = nil
}

func (srv *Server) register(c *Client) {
	srv.info.register(c)
}

func (srv *Server) deregister(clientID uint64) {
	srv.info.deregister(clientID)
	// go func() {
	// 	srv.mu.Lock()
	// 	defer srv.mu.Unlock()
	//
	// 	if srv.released != nil {
	// 		srv.released.Done()
	// 	}
	// }()
}

// Starts a new session, serving client
func (srv *Server) serveClient(c *Client, synch bool) error {
	// Register client
	srv.register(c)
	defer srv.deregister(c.id)

	var done sync.WaitGroup
	chanErr := make(chan error, 2)

	if synch {
		done.Add(1)
	} else {
		done.Add(2)
		go func(chanErr chan error) {
			if err := srv.handleResponses(c); err != nil {
				c.Close()
				chanErr <- err
			}
			done.Done()
		}(chanErr)
	}

	go func(chanErr chan error) {
		if err := srv.handleRequests(c); err != nil {
			c.Close()
			chanErr <- err
		}
		done.Done()

		// Wait all done and release client
		done.Wait()
		close(chanErr)
		c.release()
	}(chanErr)

	// Return immediately on error
	return <-chanErr
}

func (srv *Server) handleRequests(c *Client) error {
	// Create perform callback
	perform := func(name string) error {
		return srv.perform(c, name)
	}
	// Init request/response loop
	for !c.closed {
		// set deadline
		if d := srv.config.Timeout; d > 0 {
			c.cn.SetDeadline(time.Now().Add(d))
		}

		// perform pipeline
		if err := c.pipeline(perform); err != nil {
			if err == io.EOF {
				return err
			}
			c.wr.AppendError("ERR " + err.Error())

			if !resp.IsProtocolError(err) {
				err := c.wr.Flush()
				if err != nil {
					return err
				}
			}
		}

		// flush buffer, return on errors
		if err := c.wr.Flush(); err != nil {
			return err
		}
	}
	return nil
}

func (srv *Server) handleResponses(c *Client) error {
	// All response should be handled.
	for response := range c.responses {
		// find handler
		srv.mu.RLock()
		cb, ok := srv.cmds[ASYNC_CB_NAME]
		srv.mu.RUnlock()

		if !ok {
			c.wr.AppendError(UnknownCommand(ASYNC_CB_NAME))
			// flush buffer, return on errors
			if err := c.wr.Flush(); err != nil {
				return err
			}
			continue
		}

		if contextable, ok := response.(Contextable); ok {
			contextable.SetContext(context.WithValue(contextable.Context(), ctxKeyClient{}, c))
		}
		cb.(Callback).ServeCallback(c.wr, response)

		// flush buffer, return on errors
		if err := c.wr.Flush(); err != nil {
			return err
		}
	}
	return nil
}

func (srv *Server) perform(c *Client, name string) (err error) {
	norm := strings.ToLower(name)

	// find handler
	srv.mu.RLock()
	h, ok := srv.cmds[norm]
	srv.mu.RUnlock()

	if !ok {
		c.wr.AppendError(UnknownCommand(name))
		_ = c.rd.SkipCmd()
		return
	}

	// register call
	srv.info.command(c.id, norm)

	switch handler := h.(type) {
	case Handler:
		if c.cmd, err = c.readCmd(c.cmd); err != nil {
			return
		}

		handler.ServeRedeo(c.wr, c.cmd)

	case StreamHandler:
		if c.scmd, err = c.streamCmd(c.scmd); err != nil {
			return
		}
		defer c.scmd.Discard()

		handler.ServeRedeoStream(c.wr, c.scmd)
	}

	// flush when buffer is large enough
	if n := c.wr.Buffered(); n > resp.MaxBufferSize/2 {
		err = c.wr.Flush()
	}
	return
}

/*
 * Lambda store part
 */

// Lambda facing serve client
func (srv *Server) ServeForeignClient(cn net.Conn, syncs ...bool) error {
	sync := true
	if len(syncs) > 0 {
		sync = syncs[0]
	}
	return srv.serveClient(newClient(cn), sync)
}

// Lambda facing serve client
func (srv *Server) ServeClient(cli *Client, syncs ...bool) error {
	sync := true
	if len(syncs) > 0 {
		sync = syncs[0]
	}
	cli.reset(cli.cn)
	return srv.serveClient(cli, sync)
}