github.com/grafana/pyroscope@v1.18.0/pkg/util/inflight_requets.go (about)

     1  package util
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // InflightRequests utility emerged due to the need to handle request draining
     8  // at the service level.
     9  //
    10  // Ideally, this should be the responsibility of the server using the service.
    11  // However, since the server is a dependency of the service and is only shut
    12  // down after the service is stopped, requests may still arrive after the Stop
    13  // call. This issue arises from how we initialize modules.
    14  //
    15  // In other scenarios, request draining could be managed at a higher level,
    16  // such as in a load balancer or service discovery mechanism. The goal would
    17  // be to stop routing requests to an instance that is about to shut down.
    18  //
    19  // In our case, service instances that are not directly exposed to the outside
    20  // world but are discoverable via e.g, ring, kubernetes, or DNS. There's no a
    21  // _reliable_ mechanism to ensure that all the clients are aware of fact that
    22  // the instance is leaving, so requests may continue to arrive within a short
    23  // period of time. InflightRequests ensure that such requests will be rejected.
    24  type InflightRequests struct {
    25  	mu      sync.RWMutex
    26  	wg      sync.WaitGroup
    27  	allowed bool
    28  }
    29  
    30  // Open allows new requests.
    31  func (r *InflightRequests) Open() {
    32  	r.mu.Lock()
    33  	r.allowed = true
    34  	r.mu.Unlock()
    35  }
    36  
    37  // Add adds a new request if allowed.
    38  func (r *InflightRequests) Add() bool {
    39  	r.mu.RLock()
    40  	defer r.mu.RUnlock()
    41  	if !r.allowed {
    42  		return false
    43  	}
    44  	r.wg.Add(1)
    45  	return true
    46  }
    47  
    48  // Done completes a request.
    49  func (r *InflightRequests) Done() {
    50  	r.wg.Done()
    51  }
    52  
    53  // Drain prevents new requests from being accepted
    54  // and waits for all ongoing requests to complete.
    55  func (r *InflightRequests) Drain() {
    56  	r.mu.Lock()
    57  	r.allowed = false
    58  	r.mu.Unlock()
    59  	r.wg.Wait()
    60  }