github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/lb/roundtripper.go (about)

     1  package lb
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"net"
     7  	"net/http"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/golang/groupcache/singleflight"
    12  )
    13  
    14  // ErrNoFallbackNodeFound happens when the fallback routine does not manage to
    15  // find a TCP reachable node in alternative to the chosen one.
    16  var ErrNoFallbackNodeFound = errors.New("no fallback node found - whole cluster seems offline")
    17  
    18  // FallbackRoundTripper implements http.RoundTripper in a way that when an
    19  // outgoing request does not manage to succeed with its original target host,
    20  // it fallsback to a list of alternative hosts. Internally it keeps a list of
    21  // dead hosts, and pings them until they are back online, diverting traffic
    22  // back to them. This is meant to be used by ConsistentHashReverseProxy().
    23  type FallbackRoundTripper struct {
    24  	nodes []string
    25  	sf    singleflight.Group
    26  
    27  	mu sync.Mutex
    28  	// a list of offline servers that must be rechecked to see when they
    29  	// get back online. If a server is in this list, it must have a fallback
    30  	// available to which requests are sent.
    31  	fallback map[string]string
    32  }
    33  
    34  // NewRoundTripper creates a new FallbackRoundTripper and triggers the internal
    35  // host TCP health checks.
    36  func NewRoundTripper(ctx context.Context, nodes []string) *FallbackRoundTripper {
    37  	frt := &FallbackRoundTripper{
    38  		nodes:    nodes,
    39  		fallback: make(map[string]string),
    40  	}
    41  	go frt.checkHealth(ctx)
    42  	return frt
    43  }
    44  
    45  func (f *FallbackRoundTripper) checkHealth(ctx context.Context) {
    46  	tick := time.NewTicker(1 * time.Second)
    47  	defer tick.Stop()
    48  	for {
    49  		select {
    50  		case <-ctx.Done():
    51  			return
    52  		case <-tick.C:
    53  			f.mu.Lock()
    54  			if len(f.fallback) == 0 {
    55  				f.mu.Unlock()
    56  				continue
    57  			}
    58  			fallback := make(map[string]string)
    59  			for k, v := range f.fallback {
    60  				fallback[k] = v
    61  			}
    62  			f.mu.Unlock()
    63  
    64  			updatedlist := make(map[string]string)
    65  			for host, target := range fallback {
    66  				if !f.ping(host) {
    67  					updatedlist[host] = target
    68  				}
    69  			}
    70  
    71  			f.mu.Lock()
    72  			f.fallback = make(map[string]string)
    73  			for k, v := range updatedlist {
    74  				f.fallback[k] = v
    75  			}
    76  			f.mu.Unlock()
    77  		}
    78  	}
    79  }
    80  
    81  func (f *FallbackRoundTripper) ping(host string) bool {
    82  	conn, err := net.Dial("tcp", host)
    83  	if err != nil {
    84  		return false
    85  	}
    86  	conn.Close()
    87  	return true
    88  }
    89  
    90  func (f *FallbackRoundTripper) fallbackHost(targetHost, failedFallback string) string {
    91  	detected, err := f.sf.Do(targetHost, func() (interface{}, error) {
    92  		for _, node := range f.nodes {
    93  			if node != targetHost && node != failedFallback && f.ping(node) {
    94  				f.mu.Lock()
    95  				f.fallback[targetHost] = node
    96  				f.mu.Unlock()
    97  				return node, nil
    98  			}
    99  		}
   100  		return "", ErrNoFallbackNodeFound
   101  	})
   102  
   103  	if err != nil {
   104  		return ""
   105  	}
   106  	return detected.(string)
   107  
   108  }
   109  
   110  // RoundTrip implements http.RoundTrip. It tried to fullfil an *http.Request to
   111  // its original target host, falling back to a list of nodes in case of failure.
   112  // After the first failure, it consistently delivers traffic to the fallback
   113  // host, until original host is back online. If no fallback node is available,
   114  // it fails with ErrNoFallbackNodeFound. In case of cascaded failure, that is,
   115  // the fallback node is also offline, it will look for another online host.
   116  func (f *FallbackRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
   117  	targetHost := req.URL.Host
   118  
   119  	f.mu.Lock()
   120  	fallback, ok := f.fallback[targetHost]
   121  	f.mu.Unlock()
   122  	if ok {
   123  		req.URL.Host = fallback
   124  		resp, err := f.callNode(req)
   125  		if err == nil {
   126  			return resp, err
   127  		}
   128  		fallback := f.fallbackHost(targetHost, fallback)
   129  		if fallback == "" {
   130  			return nil, ErrNoFallbackNodeFound
   131  		}
   132  		req.URL.Host = fallback
   133  		return f.callNode(req)
   134  	}
   135  
   136  	resp, err := f.callNode(req)
   137  	if err == nil {
   138  		return resp, err
   139  	}
   140  
   141  	fallback = f.fallbackHost(targetHost, "")
   142  	if fallback == "" {
   143  		return nil, ErrNoFallbackNodeFound
   144  	}
   145  	req.URL.Host = fallback
   146  	return f.callNode(req)
   147  }
   148  
   149  func (f *FallbackRoundTripper) callNode(req *http.Request) (*http.Response, error) {
   150  	requestURI := req.RequestURI
   151  	req.RequestURI = ""
   152  	resp, err := http.DefaultClient.Do(req)
   153  	if err == nil {
   154  		resp.Request.RequestURI = requestURI
   155  	}
   156  	return resp, err
   157  }