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

     1  // Copyright 2022 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  	"github.com/grailbio/base/errors"
    12  )
    13  
    14  // Mutex is a context-aware mutex.  It must not be copied.
    15  // The zero value is ready to use.
    16  type Mutex struct {
    17  	initOnce sync.Once
    18  	lockCh   chan struct{}
    19  }
    20  
    21  // Lock attempts to exclusively lock m.  If the m is already locked, it will
    22  // wait until it is unlocked.  If ctx is canceled before the lock can be taken,
    23  // Lock will not take the lock, and a non-nil error is returned.
    24  func (m *Mutex) Lock(ctx context.Context) error {
    25  	m.init()
    26  	select {
    27  	case m.lockCh <- struct{}{}:
    28  		return nil
    29  	case <-ctx.Done():
    30  		return errors.E(ctx.Err(), "waiting for lock")
    31  	}
    32  }
    33  
    34  // Unlock unlocks m.  It must be called exactly once iff Lock returns nil.
    35  // Unlock panics if it is called while m is not locked.
    36  func (m *Mutex) Unlock() {
    37  	m.init()
    38  	select {
    39  	case <-m.lockCh:
    40  	default:
    41  		panic("Unlock called on mutex that is not locked")
    42  	}
    43  }
    44  
    45  func (m *Mutex) init() {
    46  	m.initOnce.Do(func() {
    47  		m.lockCh = make(chan struct{}, 1)
    48  	})
    49  }