gate.go

     1  // Copyright 2024 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  // Package gate contains an alternative condition variable.
     6  package gate
     7  
     8  import "context"
     9  
    10  // A Gate is a monitor (mutex + condition variable) with one bit of state.
    11  //
    12  // The condition may be either set or unset.
    13  // Lock operations may be unconditional, or wait for the condition to be set.
    14  // Unlock operations record the new state of the condition.
    15  type Gate struct {
    16  	// When unlocked, exactly one of set or unset contains a value.
    17  	// When locked, neither chan contains a value.
    18  	set   chan struct{}
    19  	unset chan struct{}
    20  }
    21  
    22  // New returns a new, unlocked gate.
    23  func New(set bool) Gate {
    24  	g := Gate{
    25  		set:   make(chan struct{}, 1),
    26  		unset: make(chan struct{}, 1),
    27  	}
    28  	g.Unlock(set)
    29  	return g
    30  }
    31  
    32  // Lock acquires the gate unconditionally.
    33  // It reports whether the condition is set.
    34  func (g *Gate) Lock() (set bool) {
    35  	select {
    36  	case <-g.set:
    37  		return true
    38  	case <-g.unset:
    39  		return false
    40  	}
    41  }
    42  
    43  // WaitAndLock waits until the condition is set before acquiring the gate.
    44  // If the context expires, WaitAndLock returns an error and does not acquire the gate.
    45  func (g *Gate) WaitAndLock(ctx context.Context) error {
    46  	select {
    47  	case <-g.set:
    48  		return nil
    49  	default:
    50  	}
    51  	select {
    52  	case <-g.set:
    53  		return nil
    54  	case <-ctx.Done():
    55  		return ctx.Err()
    56  	}
    57  }
    58  
    59  // LockIfSet acquires the gate if and only if the condition is set.
    60  func (g *Gate) LockIfSet() (acquired bool) {
    61  	select {
    62  	case <-g.set:
    63  		return true
    64  	default:
    65  		return false
    66  	}
    67  }
    68  
    69  // Unlock sets the condition and releases the gate.
    70  func (g *Gate) Unlock(set bool) {
    71  	if set {
    72  		g.set <- struct{}{}
    73  	} else {
    74  		g.unset <- struct{}{}
    75  	}
    76  }
    77
View as plain text