func (s *adminServer) Drain(ctx context.Context, req *serverpb.DrainRequest) (*serverpb.DrainResponse, error) {
on := make([]serverpb.DrainMode, len(req.On))
for i := range req.On {
on[i] = serverpb.DrainMode(req.On[i])
}
off := make([]serverpb.DrainMode, len(req.Off))
for i := range req.Off {
off[i] = serverpb.DrainMode(req.Off[i])
}
_ = s.server.Undrain(off)
nowOn, err := s.server.Drain(on)
if err != nil {
return nil, err
}
nowOnInts := make([]int32, len(nowOn))
for i := range nowOn {
nowOnInts[i] = int32(nowOn[i])
}
if req.Shutdown {
s.server.stopper.Quiesce()
go func() {
time.Sleep(50 * time.Millisecond)
s.server.stopper.Stop()
}()
}
return &serverpb.DrainResponse{On: nowOnInts}, nil
}
func (s *adminServer) Drain(req *serverpb.DrainRequest, stream serverpb.Admin_DrainServer) error {
on := make([]serverpb.DrainMode, len(req.On))
for i := range req.On {
on[i] = serverpb.DrainMode(req.On[i])
}
off := make([]serverpb.DrainMode, len(req.Off))
for i := range req.Off {
off[i] = serverpb.DrainMode(req.Off[i])
}
_ = s.server.Undrain(off)
nowOn, err := s.server.Drain(on)
if err != nil {
return err
}
res := serverpb.DrainResponse{
On: make([]int32, len(nowOn)),
}
for i := range nowOn {
res.On[i] = int32(nowOn[i])
}
if err := stream.Send(&res); err != nil {
return err
}
if req.Shutdown {
s.server.grpc.Stop()
go s.server.stopper.Stop()
}
ctx := stream.Context()
select {
case <-s.server.stopper.IsStopped():
return nil
case <-ctx.Done():
return ctx.Err()
}
}