trpc.group/trpc-go/trpc-go@v1.0.3/healthcheck/health_check.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package healthcheck is used to check service status.
    15  package healthcheck
    16  
    17  import (
    18  	"fmt"
    19  	"sync"
    20  )
    21  
    22  // Status is the status of server or individual services.
    23  type Status int
    24  
    25  const (
    26  	// Unknown is the initial status of a service.
    27  	Unknown Status = iota
    28  	// Serving indicate the status of service is ok.
    29  	Serving
    30  	// NotServing indicate the service is not available now.
    31  	NotServing
    32  )
    33  
    34  // New creates a new HealthCheck.
    35  func New(opts ...Opt) *HealthCheck {
    36  	hc := HealthCheck{
    37  		unregisteredServiceStatus: Unknown,
    38  		serviceStatuses:           make(map[string]Status),
    39  		statusServices: map[Status]map[string]struct{}{
    40  			Unknown:    make(map[string]struct{}),
    41  			Serving:    make(map[string]struct{}),
    42  			NotServing: make(map[string]struct{}),
    43  		},
    44  		serviceWatchers: make(map[string][]func(status Status)),
    45  	}
    46  	for _, opt := range opts {
    47  		opt(&hc)
    48  	}
    49  	return &hc
    50  }
    51  
    52  // HealthCheck is the struct to implement health check.
    53  type HealthCheck struct {
    54  	unregisteredServiceStatus Status
    55  
    56  	rwm             sync.RWMutex
    57  	serviceStatuses map[string]Status
    58  	statusServices  map[Status]map[string]struct{}
    59  	serviceWatchers map[string][]func(status Status)
    60  }
    61  
    62  // Register registers a new service with initial status Unknown and returns a function to update status.
    63  func (hc *HealthCheck) Register(name string) (update func(Status), err error) {
    64  	hc.rwm.Lock()
    65  	defer hc.rwm.Unlock()
    66  	if _, ok := hc.serviceStatuses[name]; ok {
    67  		return nil, fmt.Errorf("service %s has been registered", name)
    68  	}
    69  	hc.serviceStatuses[name] = Unknown
    70  	hc.statusServices[Unknown][name] = struct{}{}
    71  	for _, onStatusChanged := range hc.serviceWatchers[name] {
    72  		onStatusChanged(Unknown)
    73  	}
    74  	return func(status Status) {
    75  		hc.rwm.Lock()
    76  		defer hc.rwm.Unlock()
    77  		delete(hc.statusServices[hc.serviceStatuses[name]], name)
    78  		hc.statusServices[status][name] = struct{}{}
    79  		hc.serviceStatuses[name] = status
    80  
    81  		for _, onStatusChanged := range hc.serviceWatchers[name] {
    82  			onStatusChanged(status)
    83  		}
    84  	}, nil
    85  }
    86  
    87  // Unregister unregisters the registered service.
    88  func (hc *HealthCheck) Unregister(name string) {
    89  	hc.rwm.Lock()
    90  	defer hc.rwm.Unlock()
    91  	delete(hc.statusServices[hc.serviceStatuses[name]], name)
    92  	delete(hc.serviceStatuses, name)
    93  }
    94  
    95  // CheckService returns the status of a service.
    96  func (hc *HealthCheck) CheckService(name string) Status {
    97  	hc.rwm.RLock()
    98  	defer hc.rwm.RUnlock()
    99  	status, ok := hc.serviceStatuses[name]
   100  	if !ok {
   101  		return hc.unregisteredServiceStatus
   102  	}
   103  	return status
   104  }
   105  
   106  // CheckServer returns the status of the entire server.
   107  func (hc *HealthCheck) CheckServer() Status {
   108  	hc.rwm.RLock()
   109  	defer hc.rwm.RUnlock()
   110  	if len(hc.statusServices[Serving]) == len(hc.serviceStatuses) {
   111  		return Serving
   112  	}
   113  	if len(hc.statusServices[Unknown]) != 0 {
   114  		return Unknown
   115  	}
   116  	return NotServing
   117  }
   118  
   119  // Watch registers a service status watcher.
   120  func (hc *HealthCheck) Watch(serviceName string, onStatusChanged func(Status)) {
   121  	hc.rwm.Lock()
   122  	defer hc.rwm.Unlock()
   123  	hc.serviceWatchers[serviceName] = append(hc.serviceWatchers[serviceName], onStatusChanged)
   124  }