github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/file/s3file/s3transport/expiring_map.go (about)

     1  package s3transport
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/Schaudge/grailbase/file/s3file/internal/autolog"
     9  	"github.com/Schaudge/grailbase/log"
    10  )
    11  
    12  const (
    13  	// expireAfter balances saving seen IPs to distribute ongoing load vs. tying up resources
    14  	// for a long time. Given that DNS provides new S3 IP addresses every few seconds, retaining
    15  	// for an hour means I/O intensive batch jobs can maintain hundreds of S3 peers. But, an API server
    16  	// with weeks of uptime won't accrete huge numbers of old records.
    17  	expireAfter = time.Hour
    18  	// expireLoopEvery controls how frequently the expireAfter threshold is tested, so it controls
    19  	// "slack" in expireAfter. The loop takes locks that block requests, so it should not be too
    20  	// frequent (relative to request rate).
    21  	expireLoopEvery = time.Minute
    22  )
    23  
    24  type expiringMap struct {
    25  	now func() time.Time
    26  
    27  	mu sync.Mutex
    28  	// elems is URL host -> string(net.IP) -> last seen.
    29  	elems map[string]map[string]time.Time
    30  }
    31  
    32  func newExpiringMap(runPeriodic runPeriodic, now func() time.Time) *expiringMap {
    33  	s := expiringMap{now: now, elems: map[string]map[string]time.Time{}}
    34  	go runPeriodic(expireLoopEvery, s.expireOnce)
    35  	autolog.Register(s.logOnce)
    36  	return &s
    37  }
    38  
    39  func (s *expiringMap) AddAndGet(host string, newIPs []net.IP) (allIPs []net.IP) {
    40  	now := s.now()
    41  	s.mu.Lock()
    42  	defer s.mu.Unlock()
    43  	ips, ok := s.elems[host]
    44  	if !ok {
    45  		ips = map[string]time.Time{}
    46  		s.elems[host] = ips
    47  	}
    48  	for _, ip := range newIPs {
    49  		ips[string(ip)] = now
    50  	}
    51  	for ip := range ips {
    52  		allIPs = append(allIPs, net.IP(ip))
    53  	}
    54  	return
    55  }
    56  
    57  func (s *expiringMap) expireOnce(now time.Time) {
    58  	earliestUnexpiredTime := now.Add(-expireAfter)
    59  	s.mu.Lock()
    60  	for host, ips := range s.elems {
    61  		deleteBefore(ips, earliestUnexpiredTime)
    62  		if len(ips) == 0 {
    63  			delete(s.elems, host)
    64  		}
    65  	}
    66  	s.mu.Unlock()
    67  }
    68  
    69  func deleteBefore(times map[string]time.Time, threshold time.Time) {
    70  	for key, time := range times {
    71  		if time.Before(threshold) {
    72  			delete(times, key)
    73  		}
    74  	}
    75  }
    76  
    77  func (s *expiringMap) logOnce() {
    78  	s.mu.Lock()
    79  	var (
    80  		hosts          = len(s.elems)
    81  		ips, hostIPMax int
    82  	)
    83  	for _, e := range s.elems {
    84  		ips += len(e)
    85  		if len(e) > hostIPMax {
    86  			hostIPMax = len(e)
    87  		}
    88  	}
    89  	s.mu.Unlock()
    90  	log.Printf("s3file transport: hosts:%d ips:%d hostipmax:%d", hosts, ips, hostIPMax)
    91  }
    92  
    93  // runPeriodic runs the given func with the given period.
    94  type runPeriodic func(time.Duration, func(time.Time))
    95  
    96  func runPeriodicForever() runPeriodic {
    97  	return func(period time.Duration, tick func(time.Time)) {
    98  		ticker := time.NewTicker(period)
    99  		defer ticker.Stop()
   100  		for {
   101  			select {
   102  			case now := <-ticker.C:
   103  				tick(now)
   104  			}
   105  		}
   106  	}
   107  }
   108  
   109  func noOpRunPeriodic(time.Duration, func(time.Time)) {}