golang.org/x/net@v0.25.1-0.20240516223405-c87a5b62e243/quic/gate.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build go1.21
     6  
     7  package quic
     8  
     9  import "context"
    10  
    11  // An gate is a monitor (mutex + condition variable) with one bit of state.
    12  //
    13  // The condition may be either set or unset.
    14  // Lock operations may be unconditional, or wait for the condition to be set.
    15  // Unlock operations record the new state of the condition.
    16  type gate struct {
    17  	// When unlocked, exactly one of set or unset contains a value.
    18  	// When locked, neither chan contains a value.
    19  	set   chan struct{}
    20  	unset chan struct{}
    21  }
    22  
    23  // newGate returns a new, unlocked gate with the condition unset.
    24  func newGate() gate {
    25  	g := newLockedGate()
    26  	g.unlock(false)
    27  	return g
    28  }
    29  
    30  // newLocked gate returns a new, locked gate.
    31  func newLockedGate() gate {
    32  	return gate{
    33  		set:   make(chan struct{}, 1),
    34  		unset: make(chan struct{}, 1),
    35  	}
    36  }
    37  
    38  // lock acquires the gate unconditionally.
    39  // It reports whether the condition is set.
    40  func (g *gate) lock() (set bool) {
    41  	select {
    42  	case <-g.set:
    43  		return true
    44  	case <-g.unset:
    45  		return false
    46  	}
    47  }
    48  
    49  // waitAndLock waits until the condition is set before acquiring the gate.
    50  // If the context expires, waitAndLock returns an error and does not acquire the gate.
    51  func (g *gate) waitAndLock(ctx context.Context, testHooks connTestHooks) error {
    52  	if testHooks != nil {
    53  		return testHooks.waitUntil(ctx, g.lockIfSet)
    54  	}
    55  	select {
    56  	case <-g.set:
    57  		return nil
    58  	default:
    59  	}
    60  	select {
    61  	case <-g.set:
    62  		return nil
    63  	case <-ctx.Done():
    64  		return ctx.Err()
    65  	}
    66  }
    67  
    68  // lockIfSet acquires the gate if and only if the condition is set.
    69  func (g *gate) lockIfSet() (acquired bool) {
    70  	select {
    71  	case <-g.set:
    72  		return true
    73  	default:
    74  		return false
    75  	}
    76  }
    77  
    78  // unlock sets the condition and releases the gate.
    79  func (g *gate) unlock(set bool) {
    80  	if set {
    81  		g.set <- struct{}{}
    82  	} else {
    83  		g.unset <- struct{}{}
    84  	}
    85  }
    86  
    87  // unlock sets the condition to the result of f and releases the gate.
    88  // Useful in defers.
    89  func (g *gate) unlockFunc(f func() bool) {
    90  	g.unlock(f())
    91  }