144 lines
2.9 KiB
Go
144 lines
2.9 KiB
Go
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))
|
|
}
|
|
}
|
|
}
|