github.com/qiniu/go-sdk/v7@v7.27.0/storagev2/chooser/subnet_chooser.go (about)

     1  package chooser
     2  
     3  import (
     4  	"container/heap"
     5  	"context"
     6  	"net"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  type (
    12  	subnetChooser struct {
    13  		blackheap      blackheap
    14  		blackheapMutex sync.Mutex
    15  		freezeDuration time.Duration
    16  	}
    17  
    18  	// SubnetChooserConfig 子网选择器的选项
    19  	SubnetChooserConfig IPChooserConfig
    20  )
    21  
    22  // NewSubnetChooser 创建子网选择器
    23  func NewSubnetChooser(options *SubnetChooserConfig) Chooser {
    24  	if options == nil {
    25  		options = &SubnetChooserConfig{}
    26  	}
    27  	freezeDuration := options.FreezeDuration
    28  	if freezeDuration == 0 {
    29  		freezeDuration = 10 * time.Minute
    30  	}
    31  	return &subnetChooser{
    32  		blackheap: blackheap{
    33  			m:     make(map[string]*blackItem, 1024),
    34  			items: make([]*blackItem, 0, 1024),
    35  		},
    36  		freezeDuration: freezeDuration,
    37  	}
    38  }
    39  
    40  func (chooser *subnetChooser) Choose(_ context.Context, ips []net.IP, options *ChooseOptions) []net.IP {
    41  	if len(ips) == 0 {
    42  		return nil
    43  	}
    44  	if options == nil {
    45  		options = &ChooseOptions{}
    46  	}
    47  
    48  	chooser.blackheapMutex.Lock()
    49  	defer chooser.blackheapMutex.Unlock()
    50  
    51  	var chosenSubnet net.IP
    52  	chosenIPs := make([]net.IP, 0, chooser.blackheap.Len())
    53  	for _, ip := range ips {
    54  		subnetIP := makeSubnet(ip)
    55  		if len(chosenIPs) == 0 {
    56  			if item := chooser.blackheap.FindByDomainAndIp(options.Domain, subnetIP); item == nil {
    57  				chosenIPs = append(chosenIPs, ip)
    58  				chosenSubnet = subnetIP
    59  			}
    60  		} else if chosenSubnet.Equal(subnetIP) {
    61  			chosenIPs = append(chosenIPs, ip)
    62  		}
    63  	}
    64  	if len(chosenIPs) > 0 || options.FailFast {
    65  		return chosenIPs
    66  	}
    67  
    68  	var chosenExpiredAt time.Time
    69  	toFind := makeSet(ips, options.Domain, func(domain string, ip net.IP) string {
    70  		return makeMapKey(domain, makeSubnet(ip))
    71  	})
    72  	backups := make([]*blackItem, 0, chooser.blackheap.Len())
    73  	for chooser.blackheap.Len() > 0 {
    74  		firstChosen := heap.Pop(&chooser.blackheap).(*blackItem)
    75  		backups = append(backups, firstChosen)
    76  		firstChosenSubnetIP := makeSubnet(firstChosen.ip)
    77  		if _, ok := toFind[makeMapKey(firstChosen.domain, firstChosenSubnetIP)]; ok {
    78  			chosenExpiredAt = firstChosen.expiredAt
    79  			chosenSubnet = firstChosenSubnetIP
    80  			break
    81  		}
    82  	}
    83  	if chosenExpiredAt.IsZero() {
    84  		panic("chosenExpiredAt should not be empty")
    85  	}
    86  	for _, ip := range ips {
    87  		if chosenSubnet.Equal(makeSubnet(ip)) {
    88  			chosenIPs = append(chosenIPs, ip)
    89  		}
    90  	}
    91  	for _, backup := range backups {
    92  		if backup.expiredAt.After(time.Now()) {
    93  			heap.Push(&chooser.blackheap, backup)
    94  		}
    95  	}
    96  	return chosenIPs
    97  }
    98  
    99  func (chooser *subnetChooser) FeedbackGood(_ context.Context, ips []net.IP, options *FeedbackOptions) {
   100  	if len(ips) == 0 {
   101  		return
   102  	}
   103  	if options == nil {
   104  		options = &FeedbackOptions{}
   105  	}
   106  
   107  	chooser.blackheapMutex.Lock()
   108  	defer chooser.blackheapMutex.Unlock()
   109  
   110  	haveFeedback := make(map[string]struct{})
   111  	for _, ip := range ips {
   112  		subnetIP := makeSubnet(ip)
   113  		subnetIPString := subnetIP.String()
   114  		if _, ok := haveFeedback[subnetIPString]; ok {
   115  			continue
   116  		} else {
   117  			haveFeedback[subnetIPString] = struct{}{}
   118  		}
   119  		if item := chooser.blackheap.FindByDomainAndIp(options.Domain, subnetIP); item != nil {
   120  			heap.Remove(&chooser.blackheap, item.index)
   121  		}
   122  	}
   123  }
   124  
   125  func (chooser *subnetChooser) FeedbackBad(_ context.Context, ips []net.IP, options *FeedbackOptions) {
   126  	if len(ips) == 0 {
   127  		return
   128  	}
   129  	if options == nil {
   130  		options = &FeedbackOptions{}
   131  	}
   132  
   133  	chooser.blackheapMutex.Lock()
   134  	defer chooser.blackheapMutex.Unlock()
   135  
   136  	haveFeedback := make(map[string]struct{})
   137  	newExpiredAt := time.Now().Add(chooser.freezeDuration)
   138  	for _, ip := range ips {
   139  		subnetIP := makeSubnet(ip)
   140  		subnetIPString := subnetIP.String()
   141  		if _, ok := haveFeedback[subnetIPString]; ok {
   142  			continue
   143  		} else {
   144  			haveFeedback[subnetIPString] = struct{}{}
   145  		}
   146  		if item := chooser.blackheap.FindByDomainAndIp(options.Domain, subnetIP); item != nil {
   147  			if chooser.blackheap.items[item.index].expiredAt.Equal(newExpiredAt) {
   148  				continue
   149  			}
   150  			chooser.blackheap.items[item.index].expiredAt = newExpiredAt
   151  			heap.Fix(&chooser.blackheap, item.index)
   152  		} else {
   153  			heap.Push(&chooser.blackheap, &blackItem{
   154  				domain:    options.Domain,
   155  				ip:        subnetIP,
   156  				expiredAt: newExpiredAt,
   157  			})
   158  		}
   159  	}
   160  }
   161  
   162  func makeSubnet(ip net.IP) net.IP {
   163  	if ipv4 := ip.To4(); ipv4 != nil {
   164  		return ipv4.Mask(ipv4.DefaultMask())
   165  	}
   166  	return ip.Mask(net.CIDRMask(64, 128))
   167  }