github.com/grailbio/base@v0.0.11/sync/once/once.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 once contains utilities for managing actions that
     6  // must be performed exactly once.
     7  package once
     8  
     9  import (
    10  	"sync"
    11  	"sync/atomic"
    12  )
    13  
    14  // Task manages a computation that must be run at most once.
    15  // It's similar to sync.Once, except it also handles and returns errors.
    16  type Task struct {
    17  	mu   sync.Mutex
    18  	done uint32
    19  	err  error
    20  }
    21  
    22  // Do run the function do at most once. Successive invocations of Do
    23  // guarantee exactly one invocation of the function do. Do returns
    24  // the error of do's invocation.
    25  func (o *Task) Do(do func() error) error {
    26  	if atomic.LoadUint32(&o.done) == 1 {
    27  		return o.err
    28  	}
    29  	o.mu.Lock()
    30  	defer o.mu.Unlock()
    31  	if atomic.LoadUint32(&o.done) == 0 {
    32  		o.err = do()
    33  		atomic.StoreUint32(&o.done, 1)
    34  	}
    35  	return o.err
    36  }
    37  
    38  // Done returns whether the task is done.
    39  func (o *Task) Done() bool {
    40  	o.mu.Lock()
    41  	defer o.mu.Unlock()
    42  	return 1 == atomic.LoadUint32(&o.done)
    43  }
    44  
    45  // Reset resets the task effectively making it possible for `Do` to invoke the underlying do func again.
    46  // Reset will only reset the task if it was already completed.
    47  func (o *Task) Reset() {
    48  	o.mu.Lock()
    49  	defer o.mu.Unlock()
    50  	atomic.CompareAndSwapUint32(&o.done, 1, 0)
    51  }
    52  
    53  // Map coordinates actions that must happen exactly once, keyed
    54  // by user-defined keys.
    55  type Map sync.Map
    56  
    57  // Perform the provided action named by a key. Do invokes the action
    58  // exactly once for each key, and returns any errors produced by the
    59  // provided action.
    60  func (m *Map) Do(key interface{}, do func() error) error {
    61  	taskv, _ := (*sync.Map)(m).LoadOrStore(key, new(Task))
    62  	task := taskv.(*Task)
    63  	return task.Do(do)
    64  }
    65  
    66  // Forget forgets past computations associated with the provided key.
    67  func (m *Map) Forget(key interface{}) {
    68  	(*sync.Map)(m).Delete(key)
    69  }