github.com/splucs/witchcraft-go-server@v1.7.0/status/reporter/reporter.go (about)

     1  // Copyright (c) 2018 Palantir Technologies. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package reporter
    16  
    17  import (
    18  	"context"
    19  	"regexp"
    20  	"sync"
    21  
    22  	werror "github.com/palantir/witchcraft-go-error"
    23  	"github.com/palantir/witchcraft-go-server/conjure/witchcraft/api/health"
    24  	"github.com/palantir/witchcraft-go-server/status"
    25  )
    26  
    27  const slsHealthNameRegex = "^[A-Z_]+$"
    28  
    29  var _ HealthReporter = &healthReporter{}
    30  
    31  type HealthReporter interface {
    32  	status.HealthCheckSource
    33  	InitializeHealthComponent(name string) (HealthComponent, error)
    34  	GetHealthComponent(name string) (HealthComponent, bool)
    35  	UnregisterHealthComponent(name string) bool
    36  }
    37  
    38  type healthReporter struct {
    39  	mutex            sync.RWMutex
    40  	healthComponents map[health.CheckType]HealthComponent
    41  }
    42  
    43  // NewHealthReporter - creates a new HealthReporter; an implementation of status.HealthCheckSource
    44  // which initializes HealthComponents to report on the health of each individual health.CheckType.
    45  func NewHealthReporter() HealthReporter {
    46  	return newHealthReporter()
    47  }
    48  
    49  func newHealthReporter() *healthReporter {
    50  	return &healthReporter{
    51  		healthComponents: make(map[health.CheckType]HealthComponent),
    52  	}
    53  }
    54  
    55  // InitializeHealthComponent - Creates a health component for the given name where the component should be stated as
    56  // initializing until a future call modifies the initializing status. The created health component is stored in the
    57  // HealthReporter and can be fetched later by name via GetHealthComponent.
    58  // Returns ErrorState if the component name is non-SLS compliant, or the name is already in use
    59  func (r *healthReporter) InitializeHealthComponent(name string) (HealthComponent, error) {
    60  	isSLSCompliant := regexp.MustCompile(slsHealthNameRegex).MatchString
    61  	if !isSLSCompliant(name) {
    62  		return nil, werror.Error("component name is not a valid SLS health component name",
    63  			werror.SafeParam("name", name),
    64  			werror.SafeParam("validPattern", slsHealthNameRegex))
    65  	}
    66  	componentName := health.CheckType(name)
    67  	healthComponent := &healthComponent{
    68  		name:  componentName,
    69  		state: StartingState,
    70  	}
    71  
    72  	r.mutex.Lock()
    73  	defer r.mutex.Unlock()
    74  	if _, ok := r.healthComponents[componentName]; ok {
    75  		return nil, werror.Error("Health component name already exists", werror.SafeParam("name", name))
    76  	}
    77  
    78  	r.healthComponents[componentName] = healthComponent
    79  	return healthComponent, nil
    80  }
    81  
    82  // GetHealthComponent - Gets an initialized health component by name.
    83  func (r *healthReporter) GetHealthComponent(name string) (HealthComponent, bool) {
    84  	r.mutex.RLock()
    85  	defer r.mutex.RUnlock()
    86  	c, ok := r.healthComponents[health.CheckType(name)]
    87  	return c, ok
    88  }
    89  
    90  // UnregisterHealthComponent - Removes a health component by name if already initialized.
    91  func (r *healthReporter) UnregisterHealthComponent(name string) bool {
    92  	r.mutex.Lock()
    93  	defer r.mutex.Unlock()
    94  
    95  	var changesMade bool
    96  	componentName := health.CheckType(name)
    97  
    98  	if _, present := r.healthComponents[componentName]; present {
    99  		delete(r.healthComponents, componentName)
   100  		changesMade = true
   101  	}
   102  
   103  	return changesMade
   104  }
   105  
   106  // HealthStatus returns a copy of the current HealthStatus, and cannot be used to modify the current state
   107  func (r *healthReporter) HealthStatus(ctx context.Context) health.HealthStatus {
   108  	checks := make(map[health.CheckType]health.HealthCheckResult, len(r.healthComponents))
   109  	for checkType, component := range r.healthComponents {
   110  		checks[checkType] = component.GetHealthCheck()
   111  	}
   112  	return health.HealthStatus{Checks: checks}
   113  }
   114  
   115  func (r *healthReporter) getHealthCheck(check health.CheckType) (health.HealthCheckResult, bool) {
   116  	r.mutex.RLock()
   117  	defer r.mutex.RUnlock()
   118  	component, found := r.healthComponents[check]
   119  	if !found {
   120  		return health.HealthCheckResult{}, false
   121  	}
   122  	return component.GetHealthCheck(), found
   123  }