github.com/blend/go-sdk@v1.20220411.3/reverseproxy/resolver.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package reverseproxy
     9  
    10  import (
    11  	"net/http"
    12  	"sync/atomic"
    13  )
    14  
    15  // Resolver is a function that takes a request and produces a destination `url.URL`.
    16  type Resolver func(*http.Request, []*Upstream) (*Upstream, error)
    17  
    18  // RoundRobinResolver returns a closure based resolver that rotates through upstreams uniformly.
    19  func RoundRobinResolver(upstreams []*Upstream) Resolver {
    20  	if len(upstreams) == 0 {
    21  		return func(req *http.Request, upstreams []*Upstream) (*Upstream, error) {
    22  			return nil, nil
    23  		}
    24  	}
    25  
    26  	if len(upstreams) == 1 {
    27  		return func(req *http.Request, upstreams []*Upstream) (*Upstream, error) {
    28  			return upstreams[0], nil
    29  		}
    30  	}
    31  
    32  	return manyRoundRobinResolver(upstreams)
    33  }
    34  
    35  func manyRoundRobinResolver(upstreams []*Upstream) Resolver {
    36  	var index int32
    37  	total := int32(len(upstreams))
    38  	return func(req *http.Request, upstreams []*Upstream) (*Upstream, error) {
    39  		newIndex := int(atomic.AddInt32(&index, 1) % total)
    40  		return upstreams[newIndex], nil
    41  	}
    42  }