Websocket 实时拉取

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

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
}