feat 初始化

server/status.go

@@ -0,0 +1,143 @@
+package server
+
+import (
+	"context"
+	"git.hlsq.asia/alex/service-common/db/etcd"
+	"git.hlsq.asia/alex/service-common/log"
+	"go.etcd.io/etcd/api/v3/mvccpb"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"sync"
+)
+
+var (
+	statusWg       = &sync.WaitGroup{}
+	statusStopFunc context.CancelFunc
+)
+
+const (
+	serverStatusKey = "/server/status" // 状态key
+)
+
+type StatusType string
+
+const (
+	StatusInit        StatusType = "init"        // 初始化状态
+	StatusRunning     StatusType = "running"     // 运行中
+	StatusMaintenance StatusType = "maintenance" // 维护中
+	StatusWhitelist   StatusType = "whitelist"   // 白名单模式
+)
+
+var serverStatus *Status
+
+type Status struct {
+	etcdClient *etcd.Client
+	status     StatusType
+	mu         sync.RWMutex
+	callbacks  []func(status StatusType)
+}
+
+func InitStatus() {
+	serverStatus = &Status{
+		etcdClient: etcd.GetClient(),
+		status:     StatusInit,
+		mu:         sync.RWMutex{},
+	}
+}
+
+func GetServerStatus() *Status {
+	return serverStatus
+}
+
+func ListenStatus(ready *sync.WaitGroup) {
+	var stopCtx context.Context
+	stopCtx, statusStopFunc = context.WithCancel(context.Background())
+	statusWg.Add(1)
+
+	go func(stopCtx context.Context) {
+		defer statusWg.Done()
+		// 监听状态变化
+		serverStatusResp, _ := serverStatus.etcdClient.Get(serverStatusKey)
+		for _, kv := range serverStatusResp.Kvs {
+			serverStatus.onStatusChange(clientv3.EventTypePut, string(kv.Key), string(kv.Value))
+		}
+		serverStatusWatchChan := serverStatus.etcdClient.Watch(serverStatusKey,
+			clientv3.WithRev(serverStatusResp.Header.Revision+1))
+
+		ready.Done()
+		for {
+			select {
+			case msg := <-serverStatusWatchChan:
+				for _, event := range msg.Events {
+					serverStatus.onStatusChange(event.Type, string(event.Kv.Key), string(event.Kv.Value))
+				}
+			case <-stopCtx.Done():
+				return
+			}
+		}
+	}(stopCtx)
+
+	return
+}
+
+func CloseStatus() {
+	if statusStopFunc != nil {
+		statusStopFunc()
+		statusWg.Wait()
+	}
+}
+
+func (s *Status) GetStatus() StatusType {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	return s.status
+}
+
+func (s *Status) CanAccess(usn int64) bool {
+	var status StatusType
+	s.mu.RLock()
+	status = s.status
+	s.mu.RUnlock()
+	switch status {
+	case StatusRunning:
+		return true
+	case StatusMaintenance:
+		return false
+	case StatusWhitelist:
+		return GetServerWhitelist().CheckWhitelist(usn)
+	default:
+		return false
+	}
+}
+
+func (s *Status) RegisterChangeCallback(fc func(status StatusType)) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.callbacks = append(s.callbacks, fc)
+}
+
+func (s *Status) SetStatus(status string) {
+	_, err := s.etcdClient.Put(serverStatusKey, status)
+	if err != nil {
+		log.Errorf("server status set failed: %v", err)
+		return
+	}
+
+	return
+}
+
+func (s *Status) onStatusChange(t mvccpb.Event_EventType, key, value string) {
+	if key != serverStatusKey {
+		return
+	}
+
+	switch t {
+	case clientv3.EventTypePut:
+		s.mu.Lock()
+		s.status = StatusType(value)
+		s.mu.Unlock()
+		for _, fc := range s.callbacks {
+			fc(StatusType(value))
+		}
+	}
+}