simple_memcache.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"net"
	"net/textproto"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"
)

type Handler interface {
	Handle(string, io.Writer) error
}

const (
	END  = "END \n\r"
	CRLF = "\n\r"
)

var (
	MAX_LENTH  = 250
	CONN_PORT  = "11212"
	ITEM_LIMIT = 65535
	KVS        = make(map[string]string, ITEM_LIMIT)

	NOT_FOUND = []byte("NOT_FOUND \n\r")
	DELETED   = []byte("DELETED \n\r")
	STORED    = []byte("STORED \n\r")

	STATS = map[string]int{
		"cmd_get":       0,
		"cmd_set":       0,
		"get_hits":      0,
		"get_misses":    0,
		"delete_hits":   0,
		"delete_misses": 0,
		"curr_items":    0,
		"limit_items":   ITEM_LIMIT,
	}

	handlers = map[string]Handler{
		"get":     GetHandler{},
		"set":     SetHandler{},
		"delete":  DeleteHandler{},
		"stats":   StatsHandler{},
		"quit":    QuitHandler{},
		"default": DefaultHandler{},
	}
)

type CommandHandler interface {
	Handle(data string, out io.Writer)
}

type GetHandler struct{}

func (h GetHandler) Handle(line string, w io.Writer) error {
	result_string := ""
	for index, element := range strings.Fields(line) {
		if index == 0 {
			continue
		}
		value, ok := KVS[element]
		if ok == true {
			result_string = result_string + "VALUE " + element + CRLF + value + CRLF
		} else {
			STATS["get_misses"] = STATS["get_misses"] + 1
		}
	}

	result_string = result_string + "END" + CRLF
	STATS["get_hits"] = STATS["get_hits"] + 1

	STATS["cmd_get"] = STATS["get_hits"] + STATS["get_misses"]
	_, err := w.Write([]byte(result_string))
	if err != nil {
		fmt.Println("GET command failed: ", err)
	}
	return err
}

type SetHandler struct {
	pending map[string]struct{}
}

func (h SetHandler) Handle(line string, w io.Writer) error {
	// key := strings.Fields(line)[1]
	// h.pending[key] = struct{}{}
	STATS["cmd_set"] = STATS["cmd_set"] + 1
	return nil
}

type DeleteHandler struct{}

func (h DeleteHandler) Handle(line string, w io.Writer) error {
	key := strings.Fields(line)[1]
	_, ok := KVS[key]
	if ok == false {
		STATS["delete_misses"] = STATS["delete_misses"] + 1
		_, err := w.Write(NOT_FOUND)
		if err != nil {
			fmt.Println("Failed to send message to client: ", err)
		}
		return err
	} else {
		STATS["delete_hits"] = STATS["delete_hits"] + 1
		delete(KVS, key)
		STATS["curr_items"] = len(KVS)

		_, err := w.Write(DELETED)
		if err != nil {
			fmt.Println("DELETE command failed: ", err)
		}
		return err
	}
}

type StatsHandler struct{}

func (h StatsHandler) Handle(line string, w io.Writer) error {
	stats_string := ""
	for key, val := range STATS {
		fmt.Println(val)
		stats_string = stats_string + key + " " + strconv.Itoa(val) + CRLF
	}

	stats_string = stats_string + END

	_, err := w.Write([]byte(stats_string))
	if err != nil {
		fmt.Println("STATS failed: ", err)
	}
	return err
}

type QuitHandler struct{}

func (h QuitHandler) Handle(b string, w io.Writer) error {
	_, err := w.Write([]byte("Closing Connection!"))
	if err != nil {
		fmt.Println("Failed to close connection from command: ", err)
	}
	return err
}

type DefaultHandler struct{}

func (h DefaultHandler) Handle(line string, w io.Writer) error {
	key := strings.Fields(line)[1]
	value := strings.Fields(line)[0]
	KVS[key] = value

	STATS["curr_items"] = len(KVS)

	_, err := w.Write(STORED)
	if err != nil {
		fmt.Println("Failed to send message to client: ", err)
	}
	return err
}

func HandleCommand(cmd string, line string, sink io.Writer) error {
	if h, ok := handlers[cmd]; !ok {
		return fmt.Errorf("unknown command %s", cmd)
	} else {
		return h.Handle(line, sink)
	}
}

func cleanup() {
	fmt.Println("cleanup")
}

func main() {
	portPtr := flag.String("port", "11212", "server port")
	itemPtr := flag.Int("items", 65535, "items limit")
	flag.Parse()
	CONN_PORT = *portPtr
	ITEM_LIMIT = *itemPtr
	KVS = make(map[string]string, ITEM_LIMIT)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	go func() {
		<-c
		cleanup()
		os.Exit(1)
	}()

	// Listen for incoming connections
	listener, err := net.Listen("tcp", ":"+CONN_PORT)

	if err != nil {
		fmt.Println("Error starting server:", err.Error())
	}

	// Close listener when server closes
	defer listener.Close()

	fmt.Println("Simple Cache Server started on " + ":" + CONN_PORT)

	for {
		// Listen for client connections
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection: ", err.Error())
		}

		// Handle Request in Goroutine
		go handleRequest(conn)
	}
}

func handleRequest(conn net.Conn) {
	defer conn.Close()

	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	client_pending_key_to_set := ""
	ascii_regexp, _ := regexp.Compile(`[[:ascii:]]`)
L:
	for {
		line, err := tp.ReadLine()
		if err != nil {
			fmt.Println("ReadContinuedLine Failed: ", err)
			break
		}

		cmd := strings.Fields(line)[0]
		switch cmd {
		case "quit":
			HandleCommand("quit", line, conn)
			break L
		case "get":
			HandleCommand("get", line, conn)
		case "set":
			if ascii_regexp.MatchString(strings.Fields(line)[1]) != true {
				conn.Write([]byte("ERROR non-ASCII characters detected \r\n"))
				break L
			}
			if len(strings.Fields(line)[1]) > MAX_LENTH {
				conn.Write([]byte("ERROR exceeded 250 character limi \r\n"))
				break L
			}
			client_pending_key_to_set = strings.Fields(line)[1]
			HandleCommand("set", line, conn)
		case "delete":
			HandleCommand("delete", line, conn)
		case "stats":
			HandleCommand("stats", line, conn)
		default:
			fmt.Println(cmd)
			fmt.Println(line)
			fmt.Println(client_pending_key_to_set)
			newLine := line + " " + client_pending_key_to_set
			fmt.Println(newLine)
			err = HandleCommand("default", newLine, conn)
			if err != nil {
				fmt.Println("Error in parsing default value")
			} else {
				client_pending_key_to_set = ""
			}
		}
	}

}