github.com/grailbio/base@v0.0.11/sync/ctxsync/cond.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package ctxsync
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  )
    11  
    12  // A Cond is a condition variable that implements
    13  // a context-aware Wait.
    14  type Cond struct {
    15  	l     sync.Locker
    16  	waitc chan struct{}
    17  }
    18  
    19  // NewCond returns a new ContextCond based on Locker l.
    20  func NewCond(l sync.Locker) *Cond {
    21  	return &Cond{l: l}
    22  }
    23  
    24  // Broadcast notifies waiters of a state change. Broadcast must only
    25  // be called while the cond's lock is held.
    26  func (c *Cond) Broadcast() {
    27  	if c.waitc != nil {
    28  		close(c.waitc)
    29  		c.waitc = nil
    30  	}
    31  }
    32  
    33  // Done returns a channel that is closed after the next broadcast of
    34  // this Cond. Done must be called with the Cond's lock held; the lock
    35  // is released before Done returns.
    36  func (c *Cond) Done() <-chan struct{} {
    37  	if c.waitc == nil {
    38  		c.waitc = make(chan struct{})
    39  	}
    40  	waitc := c.waitc
    41  	c.l.Unlock()
    42  	return waitc
    43  }
    44  
    45  // Wait returns after the next call to Broadcast, or if the context
    46  // is complete. The context's lock must be held when calling Wait.
    47  // An error returns with the context's error if the context completes
    48  // while waiting.
    49  func (c *Cond) Wait(ctx context.Context) error {
    50  	waitc := c.Done()
    51  	var err error
    52  	select {
    53  	case <-waitc:
    54  	case <-ctx.Done():
    55  		err = ctx.Err()
    56  	}
    57  	c.l.Lock()
    58  	return err
    59  }