github.com/erda-project/erda-infra@v1.0.9/providers/health/provider.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     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 health
    16  
    17  import (
    18  	"context"
    19  	"encoding/json"
    20  	"net/http"
    21  	"reflect"
    22  	"sort"
    23  
    24  	"github.com/erda-project/erda-infra/base/servicehub"
    25  	"github.com/erda-project/erda-infra/providers/httpserver"
    26  )
    27  
    28  // Checker .
    29  type Checker func(context.Context) error
    30  
    31  // Interface .
    32  type Interface interface {
    33  	Register(Checker)
    34  }
    35  
    36  type config struct {
    37  	Path           []string `file:"path" default:"/health" desc:"http path"`
    38  	HealthStatus   int      `file:"health_status" default:"200" desc:"http response status if health"`
    39  	UnhealthStatus int      `file:"unhealth_status" default:"503" desc:"http response status if unhealth"`
    40  	HealthBody     string   `file:"health_body" desc:"http response body if health"`
    41  	UnhealthBody   string   `file:"unhealth_body" desc:"http response body if unhealth"`
    42  	ContentType    string   `file:"content_type" default:"application/json" desc:"http response Content-Type"`
    43  	AbortOnError   bool     `file:"abort_on_error"`
    44  }
    45  
    46  type provider struct {
    47  	Cfg          *config
    48  	Router       httpserver.Router `autowired:"http-server"`
    49  	names        []string
    50  	checkers     map[string][]Checker
    51  	healthBody   []byte
    52  	unhealthBody []byte
    53  }
    54  
    55  func (p *provider) Init(ctx servicehub.Context) error {
    56  	for _, path := range p.Cfg.Path {
    57  		p.Router.GET(path, p.handler)
    58  	}
    59  	p.healthBody = []byte(p.Cfg.HealthBody)
    60  	p.unhealthBody = []byte(p.Cfg.UnhealthBody)
    61  	return nil
    62  }
    63  
    64  func (p *provider) handler(resp http.ResponseWriter, req *http.Request) error {
    65  	status := make(map[string]interface{})
    66  	health := true
    67  	for _, key := range p.names {
    68  		var errors []interface{}
    69  		for _, checker := range p.checkers[key] {
    70  			err := checker(context.Background())
    71  			if err != nil {
    72  				errors = append(errors, err.Error())
    73  				health = false
    74  				if p.Cfg.AbortOnError {
    75  					break
    76  				}
    77  			}
    78  		}
    79  		status[key] = errors
    80  	}
    81  	resp.Header().Set("Content-Type", p.Cfg.ContentType)
    82  	var body []byte
    83  	if health {
    84  		resp.WriteHeader(p.Cfg.HealthStatus)
    85  		body = p.healthBody
    86  	} else {
    87  		resp.WriteHeader(p.Cfg.UnhealthStatus)
    88  		body = p.unhealthBody
    89  	}
    90  	if len(body) > 0 {
    91  		resp.Write(body)
    92  	} else {
    93  		byts, _ := json.Marshal(map[string]interface{}{
    94  			"health":   health,
    95  			"checkers": status,
    96  		})
    97  		resp.Write(byts)
    98  	}
    99  	return nil
   100  }
   101  
   102  // Provide .
   103  func (p *provider) Provide(ctx servicehub.DependencyContext, args ...interface{}) interface{} {
   104  	return &service{
   105  		name: ctx.Caller(),
   106  		p:    p,
   107  	}
   108  }
   109  
   110  type service struct {
   111  	name string
   112  	p    *provider
   113  }
   114  
   115  func (s *service) Register(c Checker) {
   116  	list, ok := s.p.checkers[s.name]
   117  	if !ok {
   118  		s.p.names = append(s.p.names, s.name)
   119  		sort.Strings(s.p.names)
   120  	}
   121  	s.p.checkers[s.name] = append(list, c)
   122  }
   123  
   124  func init() {
   125  	servicehub.Register("health", &servicehub.Spec{
   126  		Services:     []string{"health", "health-checker"},
   127  		Types:        []reflect.Type{reflect.TypeOf((*Interface)(nil)).Elem()},
   128  		Dependencies: []string{"http-server"},
   129  		Description:  "http health check",
   130  		ConfigFunc:   func() interface{} { return &config{} },
   131  		Creator: func() servicehub.Provider {
   132  			return &provider{
   133  				checkers: make(map[string][]Checker),
   134  			}
   135  		},
   136  	})
   137  }