github.com/uber/kraken@v0.1.4/lib/healthcheck/filter.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, 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  package healthcheck
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  
    20  	"github.com/uber/kraken/utils/stringset"
    21  )
    22  
    23  // Filter filters out unhealthy hosts from a host list.
    24  type Filter interface {
    25  	Run(addrs stringset.Set) stringset.Set
    26  }
    27  
    28  type filter struct {
    29  	config  FilterConfig
    30  	checker Checker
    31  	state   *state
    32  }
    33  
    34  // NewFilter creates a new Filter. Filter is stateful -- consecutive runs are required
    35  // to detect healthy / unhealthy hosts.
    36  func NewFilter(config FilterConfig, checker Checker) Filter {
    37  	config.applyDefaults()
    38  	return &filter{
    39  		config:  config,
    40  		checker: checker,
    41  		state:   newState(config),
    42  	}
    43  }
    44  
    45  // Run applies checker to addrs against the current filter state and returns the
    46  // healthy entries. New entries in addrs not found in the current state are
    47  // assumed as initially healthy. If addrs only contains a single entry, it is
    48  // always considered healthy.
    49  func (f *filter) Run(addrs stringset.Set) stringset.Set {
    50  	if len(addrs) == 1 {
    51  		return addrs.Copy()
    52  	}
    53  
    54  	f.state.sync(addrs)
    55  
    56  	ctx, cancel := context.WithTimeout(context.Background(), f.config.Timeout)
    57  	defer cancel()
    58  
    59  	var wg sync.WaitGroup
    60  	for addr := range addrs {
    61  		wg.Add(1)
    62  		go func(addr string) {
    63  			defer wg.Done()
    64  			if err := f.check(ctx, addr); err != nil {
    65  				f.state.failed(addr)
    66  			} else {
    67  				f.state.passed(addr)
    68  			}
    69  		}(addr)
    70  	}
    71  	wg.Wait()
    72  
    73  	return f.state.getHealthy()
    74  }
    75  
    76  func (f *filter) check(ctx context.Context, addr string) error {
    77  	errc := make(chan error, 1)
    78  	go func() { errc <- f.checker.Check(ctx, addr) }()
    79  	select {
    80  	case <-ctx.Done():
    81  		return ctx.Err()
    82  	case err := <-errc:
    83  		return err
    84  	}
    85  }