Buriburizaemon/nezha_domains
1
0
Fork 0
mirror of https://github.com/Buriburizaem0n/nezha_domains.git synced 2026-02-08 06:29:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Websocket 实时拉取

Browse Source
This commit is contained in:
奶爸
2019-12-10 17:57:57 +08:00
parent be57af065d
commit 3c39e1537d
11 changed files with 197 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
service/dao/dao.go
View File

@@ -7,6 +7,7 @@ import (
"github.com/patrickmn/go-cache"
"github.com/p14yground/nezha/model"
pb "github.com/p14yground/nezha/proto"
)
// Conf ..
@@ -29,3 +30,19 @@ var ServerLock sync.RWMutex
// Version ..
var Version = "debug"
// SendCommand ..
func SendCommand(cmd *pb.Command) {
ServerLock.RLock()
defer ServerLock.RUnlock()
var err error
for _, server := range ServerList {
if server.Stream != nil {
err = server.Stream.Send(cmd)
if err != nil {
close(server.StreamClose)
server.Stream = nil
}
}
}
}

22
service/rpc/nezha.go
View File

@@ -21,8 +21,8 @@ func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt,
if clientID, err = s.Auth.Check(c); err != nil {
return nil, err
}
dao.ServerLock.Lock()
defer dao.ServerLock.Unlock()
dao.ServerLock.RLock()
defer dao.ServerLock.RUnlock()
dao.ServerList[clientID].State = model.PB2State(r)
return &pb.Receipt{Proced: true}, nil
}
@@ -35,13 +35,15 @@ func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatSer
if clientID, err = s.Auth.Check(stream.Context()); err != nil {
return err
}
err = stream.Send(&pb.Command{
Type: model.MTReportState,
})
if err != nil {
log.Printf("Heartbeat stream.Send err:%v", err)
dao.ServerLock.RLock()
defer dao.ServerLock.RUnlock()
closeCh := make(chan error)
dao.ServerList[clientID].StreamClose = closeCh
dao.ServerList[clientID].Stream = stream
select {
case err = <-closeCh:
return err
}
select {}
}
// Register ..
@@ -51,8 +53,8 @@ func (s *NezhaHandler) Register(c context.Context, r *pb.Host) (*pb.Receipt, err
if clientID, err = s.Auth.Check(c); err != nil {
return nil, err
}
dao.ServerLock.Lock()
defer dao.ServerLock.Unlock()
dao.ServerLock.RLock()
defer dao.ServerLock.RUnlock()
dao.ServerList[clientID].Host = model.PB2Host(r)
return &pb.Receipt{Proced: true}, nil
}