github.com/qiuhoude/go-web@v0.0.0-20220223060959-ab545e78f20d/prepare/17epoller/redis-server/main.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"github.com/smallnest/epoller"
     6  	"github.com/smallnest/redcon"
     7  	"log"
     8  	"net"
     9  	"strconv"
    10  	"strings"
    11  	"sync"
    12  )
    13  
    14  var (
    15  	port = flag.Int("port", 6380, "server port")
    16  	pool = flag.Int("pool", 100, "pool size for epoll")
    17  )
    18  
    19  var (
    20  	mu      sync.RWMutex
    21  	keys    = make(map[string]string)
    22  	readers = make(map[net.Conn]*redcon.Reader)
    23  	writers = make(map[net.Conn]*redcon.Writer)
    24  )
    25  
    26  func main() {
    27  	flag.Parse()
    28  	ln, err := net.Listen("tcp", ":"+strconv.Itoa(*port))
    29  	if err != nil {
    30  		log.Fatalf("failed to listen %d: %v", *port, err)
    31  		return
    32  	}
    33  
    34  	poller, err := epoller.NewPollerWithBuffer(128)
    35  	if err != nil {
    36  		log.Fatalf("failed to create a poller for port %d: %v", port, err)
    37  		return
    38  	}
    39  
    40  	go poll(poller)
    41  	for {
    42  		select {
    43  		default:
    44  			conn, err := ln.Accept()
    45  			if err != nil {
    46  				return
    47  			}
    48  			poller.Add(conn)
    49  
    50  			mu.Lock()
    51  			readers[conn] = redcon.NewReader(conn)
    52  			writers[conn] = redcon.NewWriter(conn)
    53  			mu.Unlock()
    54  		}
    55  	}
    56  }
    57  
    58  func poll(poller epoller.Poller) {
    59  	for {
    60  		conns, err := poller.WaitWithBuffer()
    61  		if err != nil {
    62  			if err.Error() != "bad file descriptor" {
    63  				log.Printf("failed to poll: %v", err)
    64  			}
    65  
    66  			continue
    67  		}
    68  
    69  		for _, conn := range conns {
    70  			mu.RLock()
    71  			r := readers[conn]
    72  			w := writers[conn]
    73  			mu.RUnlock()
    74  
    75  			if r == nil || w == nil {
    76  				continue
    77  			}
    78  
    79  			cmds, err := r.ReadCommands()
    80  			if err != nil {
    81  				closeConn(conn)
    82  				continue
    83  			}
    84  
    85  			for _, cmd := range cmds {
    86  				handleCmd(conn, &cmd, w)
    87  			}
    88  		}
    89  	}
    90  }
    91  
    92  func handleCmd(conn net.Conn, cmd *redcon.Command, w *redcon.Writer) {
    93  	args := cmd.GetAllArgs()
    94  	op := strings.ToUpper(string(args[0]))
    95  	switch op {
    96  	default:
    97  		w.WriteError("ERR unknown command '" + string(args[0]) + "'")
    98  	case "PING":
    99  		if len(args) > 2 {
   100  			w.WriteError("ERR wrong number of arguments for '" + string(args[0]) + "' command")
   101  		} else if len(args) == 2 {
   102  			w.WriteBulk(args[1])
   103  		} else {
   104  			w.WriteString("PONG")
   105  		}
   106  	case "ECHO":
   107  		if len(args) != 2 {
   108  			w.WriteError("ERR wrong number of arguments for '" + string(args[0]) + "' command")
   109  		} else {
   110  			w.WriteBulk(args[1])
   111  		}
   112  	case "QUIT":
   113  		w.WriteString("OK")
   114  		w.Flush()
   115  		closeConn(conn)
   116  		return
   117  	case "GET":
   118  		if len(args) != 2 {
   119  			w.WriteError("ERR wrong number of arguments for '" + string(args[0]) + "' command")
   120  		} else {
   121  			key := string(args[1])
   122  			mu.Lock()
   123  			val, ok := keys[key]
   124  			mu.Unlock()
   125  			if !ok {
   126  				w.WriteNull()
   127  			} else {
   128  				w.WriteBulkString(val)
   129  			}
   130  		}
   131  	case "SET":
   132  		if len(args) != 3 {
   133  			w.WriteError("ERR wrong number of arguments for '" + string(args[0]) + "' command")
   134  		} else {
   135  			key, val := string(args[1]), string(args[2])
   136  			mu.Lock()
   137  			keys[key] = val
   138  			mu.Unlock()
   139  			w.WriteString("OK")
   140  		}
   141  	case "DEL":
   142  		if len(args) < 2 {
   143  			w.WriteError("ERR wrong number of arguments for '" + string(args[0]) + "' command")
   144  		} else {
   145  			var n int
   146  			mu.Lock()
   147  			for i := 1; i < len(args); i++ {
   148  				if _, ok := keys[string(args[i])]; ok {
   149  					n++
   150  					delete(keys, string(args[i]))
   151  				}
   152  			}
   153  			mu.Unlock()
   154  			w.WriteInt64(int64(n))
   155  		}
   156  	case "FLUSHDB":
   157  		mu.Lock()
   158  		keys = make(map[string]string)
   159  		mu.Unlock()
   160  		w.WriteString("OK")
   161  	}
   162  	w.Flush()
   163  }
   164  
   165  func closeConn(conn net.Conn) {
   166  	conn.Close()
   167  	mu.Lock()
   168  	delete(readers, conn)
   169  	delete(writers, conn)
   170  	mu.Unlock()
   171  }