github.com/hdt3213/godis@v1.2.9/redis/server/server.go (about)

     1  package server
     2  
     3  /*
     4   * A tcp.Handler implements redis protocol
     5   */
     6  
     7  import (
     8  	"context"
     9  	"io"
    10  	"net"
    11  	"strings"
    12  	"sync"
    13  
    14  	"github.com/hdt3213/godis/cluster"
    15  	"github.com/hdt3213/godis/config"
    16  	database2 "github.com/hdt3213/godis/database"
    17  	"github.com/hdt3213/godis/interface/database"
    18  	"github.com/hdt3213/godis/lib/logger"
    19  	"github.com/hdt3213/godis/lib/sync/atomic"
    20  	"github.com/hdt3213/godis/redis/connection"
    21  	"github.com/hdt3213/godis/redis/parser"
    22  	"github.com/hdt3213/godis/redis/protocol"
    23  )
    24  
    25  var (
    26  	unknownErrReplyBytes = []byte("-ERR unknown\r\n")
    27  )
    28  
    29  // Handler implements tcp.Handler and serves as a redis server
    30  type Handler struct {
    31  	activeConn sync.Map // *client -> placeholder
    32  	db         database.DB
    33  	closing    atomic.Boolean // refusing new client and new request
    34  }
    35  
    36  // MakeHandler creates a Handler instance
    37  func MakeHandler() *Handler {
    38  	var db database.DB
    39  	if config.Properties.Self != "" &&
    40  		len(config.Properties.Peers) > 0 {
    41  		db = cluster.MakeCluster()
    42  	} else {
    43  		db = database2.NewStandaloneServer()
    44  	}
    45  	return &Handler{
    46  		db: db,
    47  	}
    48  }
    49  
    50  func (h *Handler) closeClient(client *connection.Connection) {
    51  	_ = client.Close()
    52  	h.db.AfterClientClose(client)
    53  	h.activeConn.Delete(client)
    54  }
    55  
    56  // Handle receives and executes redis commands
    57  func (h *Handler) Handle(ctx context.Context, conn net.Conn) {
    58  	if h.closing.Get() {
    59  		// closing handler refuse new connection
    60  		_ = conn.Close()
    61  		return
    62  	}
    63  
    64  	client := connection.NewConn(conn)
    65  	h.activeConn.Store(client, struct{}{})
    66  
    67  	ch := parser.ParseStream(conn)
    68  	for payload := range ch {
    69  		if payload.Err != nil {
    70  			if payload.Err == io.EOF ||
    71  				payload.Err == io.ErrUnexpectedEOF ||
    72  				strings.Contains(payload.Err.Error(), "use of closed network connection") {
    73  				// connection closed
    74  				h.closeClient(client)
    75  				logger.Info("connection closed: " + client.RemoteAddr().String())
    76  				return
    77  			}
    78  			// protocol err
    79  			errReply := protocol.MakeErrReply(payload.Err.Error())
    80  			_, err := client.Write(errReply.ToBytes())
    81  			if err != nil {
    82  				h.closeClient(client)
    83  				logger.Info("connection closed: " + client.RemoteAddr().String())
    84  				return
    85  			}
    86  			continue
    87  		}
    88  		if payload.Data == nil {
    89  			logger.Error("empty payload")
    90  			continue
    91  		}
    92  		r, ok := payload.Data.(*protocol.MultiBulkReply)
    93  		if !ok {
    94  			logger.Error("require multi bulk protocol")
    95  			continue
    96  		}
    97  		result := h.db.Exec(client, r.Args)
    98  		if result != nil {
    99  			_, _ = client.Write(result.ToBytes())
   100  		} else {
   101  			_, _ = client.Write(unknownErrReplyBytes)
   102  		}
   103  	}
   104  }
   105  
   106  // Close stops handler
   107  func (h *Handler) Close() error {
   108  	logger.Info("handler shutting down...")
   109  	h.closing.Set(true)
   110  	// TODO: concurrent wait
   111  	h.activeConn.Range(func(key interface{}, val interface{}) bool {
   112  		client := key.(*connection.Connection)
   113  		_ = client.Close()
   114  		return true
   115  	})
   116  	h.db.Close()
   117  	return nil
   118  }