GithubBackup/httprunner
1
0
Fork 0
mirror of https://github.com/httprunner/httprunner.git synced 2026-09-05 23:51:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: state machine

Browse Source
This commit is contained in:
徐聪
2022-07-25 21:30:01 +08:00
parent 8fa04335e9
commit 76e058d709
3 changed files with 78 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hrp/internal/boomer/server_grpc.go
+32 -24
View File
@@ -138,6 +138,8 @@ type grpcServer struct {
fromWorker chan *genericMessage
disconnectedChan chan bool
shutdownChan chan bool
wg sync.WaitGroup
}
var (
@@ -221,6 +223,7 @@ func newServer(masterHost string, masterPort int) (server *grpcServer) {
fromWorker: make(chan *genericMessage, 100),
disconnectedChan: make(chan bool),
shutdownChan: make(chan bool),
wg: sync.WaitGroup{},
}
return server
}
@@ -290,6 +293,8 @@ func (s *grpcServer) valid(token string) (isValid bool) {
}
func (s *grpcServer) BidirectionalStreamingMessage(srv messager.Message_BidirectionalStreamingMessageServer) error {
s.wg.Add(1)
defer s.wg.Done()
token, ok := extractToken(srv.Context())
if !ok {
return status.Error(codes.Unauthenticated, "missing token header")
@@ -303,32 +308,34 @@ func (s *grpcServer) BidirectionalStreamingMessage(srv messager.Message_Bidirect
go s.sendMsg(srv, token)
FOR:
for {
msg, err := srv.Recv()
if st, ok := status.FromError(err); ok {
switch st.Code() {
case codes.OK:
s.fromWorker <- newGenericMessage(msg.Type, msg.Data, msg.NodeID)
log.Info().
Str("nodeID", msg.NodeID).
Str("type", msg.Type).
Interface("data", msg.Data).
Msg("receive data from worker")
case codes.Unavailable, codes.Canceled, codes.DeadlineExceeded:
s.fromWorker <- newQuitMessage(token)
break FOR
default:
log.Error().Err(err).Msg("failed to get stream from client")
break FOR
select {
case <-srv.Context().Done():
break FOR
case <-s.disconnectedChannel():
break FOR
default:
msg, err := srv.Recv()
if st, ok := status.FromError(err); ok {
switch st.Code() {
case codes.OK:
s.fromWorker <- newGenericMessage(msg.Type, msg.Data, msg.NodeID)
log.Info().
Str("nodeID", msg.NodeID).
Str("type", msg.Type).
Interface("data", msg.Data).
Msg("receive data from worker")
case codes.Unavailable, codes.Canceled, codes.DeadlineExceeded:
s.fromWorker <- newQuitMessage(token)
break FOR
default:
log.Error().Err(err).Msg("failed to get stream from client")
break FOR
}
}
}
}
// disconnected to worker
select {
case <-srv.Context().Done():
return srv.Context().Err()
case <-s.disconnectedChan:
}
log.Warn().Str("worker id", token).Msg("worker quited")
log.Info().Str("worker id", token).Msg("bidirectional stream closed")
return nil
}
@@ -338,6 +345,8 @@ func (s *grpcServer) sendMsg(srv messager.Message_BidirectionalStreamingMessageS
select {
case <-srv.Context().Done():
return
case <-s.disconnectedChannel():
return
case res := <-stream:
if s, ok := status.FromError(srv.Send(res)); ok {
switch s.Code() {
@@ -406,7 +415,6 @@ func (s *grpcServer) close() {
ctx, cancel := context.WithTimeout(context.Background(), timeout)
s.stopServer(ctx)
cancel()
close(s.disconnectedChan)
}
func (s *grpcServer) recvChannel() chan *genericMessage {