github.com/flower-corp/rosedb@v1.1.2-0.20230117132829-21dc4f7b319a/cmd/server.go (about)

     1  package main
     2  
     3  import (
     4  	"errors"
     5  	"flag"
     6  	"fmt"
     7  	"github.com/flower-corp/rosedb"
     8  	"github.com/flower-corp/rosedb/logger"
     9  	"github.com/tidwall/redcon"
    10  	"io/ioutil"
    11  	"os"
    12  	"os/signal"
    13  	"path/filepath"
    14  	"sync"
    15  	"syscall"
    16  	"time"
    17  )
    18  
    19  var (
    20  	errClientIsNil = errors.New("ERR client conn is nil")
    21  )
    22  
    23  var (
    24  	defaultDBPath            = filepath.Join("/tmp", "rosedb")
    25  	defaultHost              = "127.0.0.1"
    26  	defaultPort              = "5200"
    27  	defaultDatabasesNum uint = 16
    28  )
    29  
    30  const (
    31  	dbName = "rosedb-%04d"
    32  )
    33  
    34  func init() {
    35  	// print banner
    36  	path, _ := filepath.Abs("resource/banner.txt")
    37  	banner, _ := ioutil.ReadFile(path)
    38  	fmt.Println(string(banner))
    39  }
    40  
    41  type Server struct {
    42  	dbs    map[int]*rosedb.RoseDB
    43  	ser    *redcon.Server
    44  	signal chan os.Signal
    45  	opts   ServerOptions
    46  	mu     *sync.RWMutex
    47  }
    48  
    49  type ServerOptions struct {
    50  	dbPath    string
    51  	host      string
    52  	port      string
    53  	databases uint
    54  }
    55  
    56  func main() {
    57  	// init server options
    58  	serverOpts := new(ServerOptions)
    59  	flag.StringVar(&serverOpts.dbPath, "dbpath", defaultDBPath, "db path")
    60  	flag.StringVar(&serverOpts.host, "host", defaultHost, "server host")
    61  	flag.StringVar(&serverOpts.port, "port", defaultPort, "server port")
    62  	flag.UintVar(&serverOpts.databases, "databases", defaultDatabasesNum, "the number of databases")
    63  	flag.Parse()
    64  
    65  	// open a default database
    66  	path := filepath.Join(serverOpts.dbPath, fmt.Sprintf(dbName, 0))
    67  	opts := rosedb.DefaultOptions(path)
    68  	now := time.Now()
    69  	db, err := rosedb.Open(opts)
    70  	if err != nil {
    71  		logger.Errorf("open rosedb err, fail to start server. %v", err)
    72  		return
    73  	}
    74  	logger.Infof("open db from [%s] successfully, time cost: %v", serverOpts.dbPath, time.Since(now))
    75  
    76  	sig := make(chan os.Signal, 1)
    77  	signal.Notify(sig, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
    78  
    79  	dbs := make(map[int]*rosedb.RoseDB)
    80  	dbs[0] = db
    81  
    82  	// init and start server
    83  	svr := &Server{dbs: dbs, signal: sig, opts: *serverOpts, mu: new(sync.RWMutex)}
    84  	addr := svr.opts.host + ":" + svr.opts.port
    85  	redServer := redcon.NewServerNetwork("tcp", addr, execClientCommand, svr.redconAccept,
    86  		func(conn redcon.Conn, err error) {
    87  		},
    88  	)
    89  	svr.ser = redServer
    90  	go svr.listen()
    91  	<-svr.signal
    92  	svr.stop()
    93  }
    94  
    95  func (svr *Server) listen() {
    96  	logger.Infof("rosedb server is running, ready to accept connections")
    97  	if err := svr.ser.ListenAndServe(); err != nil {
    98  		logger.Fatalf("listen and serve err, fail to start. %v", err)
    99  		return
   100  	}
   101  }
   102  
   103  func (svr *Server) stop() {
   104  	for _, db := range svr.dbs {
   105  		if err := db.Close(); err != nil {
   106  			logger.Errorf("close db err: %v", err)
   107  		}
   108  	}
   109  	if err := svr.ser.Close(); err != nil {
   110  		logger.Errorf("close server err: %v", err)
   111  	}
   112  	logger.Infof("rosedb is ready to exit, bye bye...")
   113  }
   114  
   115  func (svr *Server) redconAccept(conn redcon.Conn) bool {
   116  	cli := new(Client)
   117  	cli.svr = svr
   118  	svr.mu.RLock()
   119  	cli.db = svr.dbs[0]
   120  	svr.mu.RUnlock()
   121  	conn.SetContext(cli)
   122  	return true
   123  }