Compare commits
4 Commits
main
...
feat_healt
Author | SHA1 | Date |
---|---|---|
haiyux | 156c97181e | 3 years ago |
haiyux | 5b6bf88a27 | 3 years ago |
haiyux | 49f700124b | 3 years ago |
haiyux | 2f74c9dfae | 3 years ago |
@ -0,0 +1,86 @@ |
|||||||
|
package health |
||||||
|
|
||||||
|
import ( |
||||||
|
"context" |
||||||
|
"sync" |
||||||
|
"time" |
||||||
|
|
||||||
|
"golang.org/x/sync/errgroup" |
||||||
|
) |
||||||
|
|
||||||
|
type Health struct { |
||||||
|
mutex sync.RWMutex |
||||||
|
// key: service name, type is string
|
||||||
|
// value: service Status type is Status
|
||||||
|
statusMap sync.Map |
||||||
|
updates map[string]map[string]chan Status |
||||||
|
ticker *time.Ticker |
||||||
|
} |
||||||
|
|
||||||
|
func New(opts ...Option) *Health { |
||||||
|
option := options{ |
||||||
|
watchTime: time.Second * 5, |
||||||
|
} |
||||||
|
for _, o := range opts { |
||||||
|
o(&option) |
||||||
|
} |
||||||
|
h := &Health{ |
||||||
|
ticker: time.NewTicker(option.watchTime), |
||||||
|
updates: map[string]map[string]chan Status{}, |
||||||
|
} |
||||||
|
return h |
||||||
|
} |
||||||
|
|
||||||
|
func (h *Health) GetStatus(service string) (Status, bool) { |
||||||
|
status, ok := h.statusMap.Load(service) |
||||||
|
if !ok { |
||||||
|
return Status_UNKNOWN, false |
||||||
|
} |
||||||
|
return status.(Status), ok |
||||||
|
} |
||||||
|
|
||||||
|
func (h *Health) SetStatus(service string, status Status) error { |
||||||
|
h.statusMap.Store(service, status) |
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100) |
||||||
|
defer cancel() |
||||||
|
eg, ctx := errgroup.WithContext(ctx) |
||||||
|
h.mutex.RLock() |
||||||
|
for _, w := range h.updates[service] { |
||||||
|
ch := w |
||||||
|
eg.Go(func() error { |
||||||
|
select { |
||||||
|
case ch <- status: |
||||||
|
case <-ctx.Done(): |
||||||
|
return ctx.Err() |
||||||
|
} |
||||||
|
return nil |
||||||
|
}) |
||||||
|
} |
||||||
|
h.mutex.RUnlock() |
||||||
|
|
||||||
|
return eg.Wait() |
||||||
|
} |
||||||
|
|
||||||
|
func (h *Health) Update(service string, id string) chan Status { |
||||||
|
h.mutex.RLock() |
||||||
|
defer h.mutex.RUnlock() |
||||||
|
if _, ok := h.updates[service]; !ok { |
||||||
|
h.updates[service] = make(map[string]chan Status) |
||||||
|
} |
||||||
|
if _, ok := h.updates[service][id]; !ok { |
||||||
|
h.updates[service][id] = make(chan Status, 1) |
||||||
|
} |
||||||
|
return h.updates[service][id] |
||||||
|
} |
||||||
|
|
||||||
|
func (h *Health) DelUpdate(service string, id string) { |
||||||
|
h.mutex.Lock() |
||||||
|
defer h.mutex.Unlock() |
||||||
|
if _, ok := h.updates[service]; ok { |
||||||
|
delete(h.updates[service], id) |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
func (h *Health) Ticker() *time.Ticker { |
||||||
|
return h.ticker |
||||||
|
} |
@ -0,0 +1,137 @@ |
|||||||
|
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||||
|
// versions:
|
||||||
|
// protoc-gen-go v1.27.1
|
||||||
|
// protoc v3.17.3
|
||||||
|
// source: health/health.proto
|
||||||
|
|
||||||
|
package health |
||||||
|
|
||||||
|
import ( |
||||||
|
protoreflect "google.golang.org/protobuf/reflect/protoreflect" |
||||||
|
protoimpl "google.golang.org/protobuf/runtime/protoimpl" |
||||||
|
reflect "reflect" |
||||||
|
sync "sync" |
||||||
|
) |
||||||
|
|
||||||
|
const ( |
||||||
|
// Verify that this generated code is sufficiently up-to-date.
|
||||||
|
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) |
||||||
|
// Verify that runtime/protoimpl is sufficiently up-to-date.
|
||||||
|
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) |
||||||
|
) |
||||||
|
|
||||||
|
type Status int32 |
||||||
|
|
||||||
|
const ( |
||||||
|
Status_UNKNOWN Status = 0 |
||||||
|
Status_SERVING Status = 1 |
||||||
|
Status_NOT_SERVING Status = 2 |
||||||
|
Status_SERVICE_UNKNOWN Status = 3 |
||||||
|
) |
||||||
|
|
||||||
|
// Enum value maps for Status.
|
||||||
|
var ( |
||||||
|
Status_name = map[int32]string{ |
||||||
|
0: "UNKNOWN", |
||||||
|
1: "SERVING", |
||||||
|
2: "NOT_SERVING", |
||||||
|
3: "SERVICE_UNKNOWN", |
||||||
|
} |
||||||
|
Status_value = map[string]int32{ |
||||||
|
"UNKNOWN": 0, |
||||||
|
"SERVING": 1, |
||||||
|
"NOT_SERVING": 2, |
||||||
|
"SERVICE_UNKNOWN": 3, |
||||||
|
} |
||||||
|
) |
||||||
|
|
||||||
|
func (x Status) Enum() *Status { |
||||||
|
p := new(Status) |
||||||
|
*p = x |
||||||
|
return p |
||||||
|
} |
||||||
|
|
||||||
|
func (x Status) String() string { |
||||||
|
return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) |
||||||
|
} |
||||||
|
|
||||||
|
func (Status) Descriptor() protoreflect.EnumDescriptor { |
||||||
|
return file_health_health_proto_enumTypes[0].Descriptor() |
||||||
|
} |
||||||
|
|
||||||
|
func (Status) Type() protoreflect.EnumType { |
||||||
|
return &file_health_health_proto_enumTypes[0] |
||||||
|
} |
||||||
|
|
||||||
|
func (x Status) Number() protoreflect.EnumNumber { |
||||||
|
return protoreflect.EnumNumber(x) |
||||||
|
} |
||||||
|
|
||||||
|
// Deprecated: Use Status.Descriptor instead.
|
||||||
|
func (Status) EnumDescriptor() ([]byte, []int) { |
||||||
|
return file_health_health_proto_rawDescGZIP(), []int{0} |
||||||
|
} |
||||||
|
|
||||||
|
var File_health_health_proto protoreflect.FileDescriptor |
||||||
|
|
||||||
|
var file_health_health_proto_rawDesc = []byte{ |
||||||
|
0x0a, 0x13, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x2f, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x2e, |
||||||
|
0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x06, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x2a, 0x48, 0x0a, |
||||||
|
0x06, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x0b, 0x0a, 0x07, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, |
||||||
|
0x57, 0x4e, 0x10, 0x00, 0x12, 0x0b, 0x0a, 0x07, 0x53, 0x45, 0x52, 0x56, 0x49, 0x4e, 0x47, 0x10, |
||||||
|
0x01, 0x12, 0x0f, 0x0a, 0x0b, 0x4e, 0x4f, 0x54, 0x5f, 0x53, 0x45, 0x52, 0x56, 0x49, 0x4e, 0x47, |
||||||
|
0x10, 0x02, 0x12, 0x13, 0x0a, 0x0f, 0x53, 0x45, 0x52, 0x56, 0x49, 0x43, 0x45, 0x5f, 0x55, 0x4e, |
||||||
|
0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x03, 0x42, 0x2e, 0x5a, 0x2c, 0x67, 0x69, 0x74, 0x68, 0x75, |
||||||
|
0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x67, 0x6f, 0x2d, 0x6b, 0x72, 0x61, 0x74, 0x6f, 0x73, 0x2f, |
||||||
|
0x6b, 0x72, 0x61, 0x74, 0x6f, 0x73, 0x2f, 0x76, 0x32, 0x2f, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, |
||||||
|
0x3b, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, |
||||||
|
} |
||||||
|
|
||||||
|
var ( |
||||||
|
file_health_health_proto_rawDescOnce sync.Once |
||||||
|
file_health_health_proto_rawDescData = file_health_health_proto_rawDesc |
||||||
|
) |
||||||
|
|
||||||
|
func file_health_health_proto_rawDescGZIP() []byte { |
||||||
|
file_health_health_proto_rawDescOnce.Do(func() { |
||||||
|
file_health_health_proto_rawDescData = protoimpl.X.CompressGZIP(file_health_health_proto_rawDescData) |
||||||
|
}) |
||||||
|
return file_health_health_proto_rawDescData |
||||||
|
} |
||||||
|
|
||||||
|
var file_health_health_proto_enumTypes = make([]protoimpl.EnumInfo, 1) |
||||||
|
var file_health_health_proto_goTypes = []interface{}{ |
||||||
|
(Status)(0), // 0: health.Status
|
||||||
|
} |
||||||
|
var file_health_health_proto_depIdxs = []int32{ |
||||||
|
0, // [0:0] is the sub-list for method output_type
|
||||||
|
0, // [0:0] is the sub-list for method input_type
|
||||||
|
0, // [0:0] is the sub-list for extension type_name
|
||||||
|
0, // [0:0] is the sub-list for extension extendee
|
||||||
|
0, // [0:0] is the sub-list for field type_name
|
||||||
|
} |
||||||
|
|
||||||
|
func init() { file_health_health_proto_init() } |
||||||
|
func file_health_health_proto_init() { |
||||||
|
if File_health_health_proto != nil { |
||||||
|
return |
||||||
|
} |
||||||
|
type x struct{} |
||||||
|
out := protoimpl.TypeBuilder{ |
||||||
|
File: protoimpl.DescBuilder{ |
||||||
|
GoPackagePath: reflect.TypeOf(x{}).PkgPath(), |
||||||
|
RawDescriptor: file_health_health_proto_rawDesc, |
||||||
|
NumEnums: 1, |
||||||
|
NumMessages: 0, |
||||||
|
NumExtensions: 0, |
||||||
|
NumServices: 0, |
||||||
|
}, |
||||||
|
GoTypes: file_health_health_proto_goTypes, |
||||||
|
DependencyIndexes: file_health_health_proto_depIdxs, |
||||||
|
EnumInfos: file_health_health_proto_enumTypes, |
||||||
|
}.Build() |
||||||
|
File_health_health_proto = out.File |
||||||
|
file_health_health_proto_rawDesc = nil |
||||||
|
file_health_health_proto_goTypes = nil |
||||||
|
file_health_health_proto_depIdxs = nil |
||||||
|
} |
@ -0,0 +1,12 @@ |
|||||||
|
syntax = "proto3"; |
||||||
|
|
||||||
|
package health; |
||||||
|
|
||||||
|
option go_package = "github.com/go-kratos/kratos/v2/health;health"; |
||||||
|
|
||||||
|
enum Status { |
||||||
|
UNKNOWN = 0; |
||||||
|
SERVING = 1; |
||||||
|
NOT_SERVING = 2; |
||||||
|
SERVICE_UNKNOWN = 3; |
||||||
|
} |
@ -0,0 +1,15 @@ |
|||||||
|
package health |
||||||
|
|
||||||
|
import "time" |
||||||
|
|
||||||
|
type Option func(*options) |
||||||
|
|
||||||
|
type options struct { |
||||||
|
watchTime time.Duration |
||||||
|
} |
||||||
|
|
||||||
|
func WithWatchTime(t time.Duration) Option { |
||||||
|
return func(o *options) { |
||||||
|
o.watchTime = t |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,99 @@ |
|||||||
|
package health |
||||||
|
|
||||||
|
import ( |
||||||
|
"context" |
||||||
|
|
||||||
|
"github.com/go-kratos/kratos/v2" |
||||||
|
"github.com/go-kratos/kratos/v2/errors" |
||||||
|
"github.com/go-kratos/kratos/v2/health" |
||||||
|
"github.com/google/uuid" |
||||||
|
pb "google.golang.org/grpc/health/grpc_health_v1" |
||||||
|
) |
||||||
|
|
||||||
|
type Server struct { |
||||||
|
pb.UnimplementedHealthServer |
||||||
|
} |
||||||
|
|
||||||
|
func NewServer() *Server { |
||||||
|
return &Server{} |
||||||
|
} |
||||||
|
|
||||||
|
func (s *Server) Check(ctx context.Context, req *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error) { |
||||||
|
info, ok := kratos.FromContext(ctx) |
||||||
|
if !ok { |
||||||
|
return nil, errors.InternalServer("kratos.FromContext(ctx) failed", "no info found in context") |
||||||
|
} |
||||||
|
status, ok := info.Health().GetStatus(req.Service) |
||||||
|
if !ok { |
||||||
|
status = health.Status_UNKNOWN |
||||||
|
} |
||||||
|
var sv pb.HealthCheckResponse_ServingStatus |
||||||
|
switch status { |
||||||
|
case health.Status_SERVING: |
||||||
|
sv = pb.HealthCheckResponse_SERVING |
||||||
|
case health.Status_NOT_SERVING: |
||||||
|
sv = pb.HealthCheckResponse_NOT_SERVING |
||||||
|
case health.Status_SERVICE_UNKNOWN: |
||||||
|
sv = pb.HealthCheckResponse_SERVICE_UNKNOWN |
||||||
|
default: |
||||||
|
sv = pb.HealthCheckResponse_NOT_SERVING |
||||||
|
} |
||||||
|
return &pb.HealthCheckResponse{ |
||||||
|
Status: sv, |
||||||
|
}, nil |
||||||
|
} |
||||||
|
|
||||||
|
func (s *Server) Watch(req *pb.HealthCheckRequest, ss pb.Health_WatchServer) (err error) { |
||||||
|
ctx := ss.Context() |
||||||
|
info, ok := kratos.FromContext(ctx) |
||||||
|
if !ok { |
||||||
|
return errors.InternalServer("get info failed", "") |
||||||
|
} |
||||||
|
uid, err := uuid.NewUUID() |
||||||
|
if err != nil { |
||||||
|
return errors.InternalServer("new uuid failed", err.Error()) |
||||||
|
} |
||||||
|
update := info.Health().Update(req.Service, uid.String()) |
||||||
|
defer info.Health().DelUpdate(req.Service, uid.String()) |
||||||
|
status, ok := info.Health().GetStatus(req.Service) |
||||||
|
if !ok { |
||||||
|
update <- health.Status_SERVICE_UNKNOWN |
||||||
|
} else { |
||||||
|
update <- status |
||||||
|
} |
||||||
|
ticker := info.Health().Ticker() |
||||||
|
for { |
||||||
|
select { |
||||||
|
case <-ctx.Done(): |
||||||
|
return nil |
||||||
|
case status = <-update: |
||||||
|
err := send(ss, status) |
||||||
|
if err != nil { |
||||||
|
return err |
||||||
|
} |
||||||
|
case <-ticker.C: |
||||||
|
err := send(ss, status) |
||||||
|
if err != nil { |
||||||
|
return err |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
func send(ss pb.Health_WatchServer, status health.Status) error { |
||||||
|
var sv pb.HealthCheckResponse_ServingStatus |
||||||
|
switch status { |
||||||
|
case health.Status_SERVING: |
||||||
|
sv = pb.HealthCheckResponse_SERVING |
||||||
|
case health.Status_NOT_SERVING: |
||||||
|
sv = pb.HealthCheckResponse_NOT_SERVING |
||||||
|
case health.Status_SERVICE_UNKNOWN: |
||||||
|
sv = pb.HealthCheckResponse_SERVICE_UNKNOWN |
||||||
|
default: |
||||||
|
sv = pb.HealthCheckResponse_NOT_SERVING |
||||||
|
} |
||||||
|
reply := &pb.HealthCheckResponse{ |
||||||
|
Status: sv, |
||||||
|
} |
||||||
|
return ss.Send(reply) |
||||||
|
} |
Loading…
Reference in new issue