github.com/linchen2chris/hugo@v0.0.0-20230307053224-cec209389705/lazy/once.go (about)

     1  // Copyright 2019 The Hugo Authors. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  // http://www.apache.org/licenses/LICENSE-2.0
     7  //
     8  // Unless required by applicable law or agreed to in writing, software
     9  // distributed under the License is distributed on an "AS IS" BASIS,
    10  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package lazy
    15  
    16  import (
    17  	"sync"
    18  	"sync/atomic"
    19  )
    20  
    21  // onceMore is similar to sync.Once.
    22  //
    23  // Additional features are:
    24  // * it can be reset, so the action can be repeated if needed
    25  // * it has methods to check if it's done or in progress
    26  //
    27  type onceMore struct {
    28  	mu   sync.Mutex
    29  	lock uint32
    30  	done uint32
    31  }
    32  
    33  func (t *onceMore) Do(f func()) {
    34  	if atomic.LoadUint32(&t.done) == 1 {
    35  		return
    36  	}
    37  
    38  	// f may call this Do and we would get a deadlock.
    39  	locked := atomic.CompareAndSwapUint32(&t.lock, 0, 1)
    40  	if !locked {
    41  		return
    42  	}
    43  	defer atomic.StoreUint32(&t.lock, 0)
    44  
    45  	t.mu.Lock()
    46  	defer t.mu.Unlock()
    47  
    48  	// Double check
    49  	if t.done == 1 {
    50  		return
    51  	}
    52  	defer atomic.StoreUint32(&t.done, 1)
    53  	f()
    54  }
    55  
    56  func (t *onceMore) InProgress() bool {
    57  	return atomic.LoadUint32(&t.lock) == 1
    58  }
    59  
    60  func (t *onceMore) Done() bool {
    61  	return atomic.LoadUint32(&t.done) == 1
    62  }
    63  
    64  func (t *onceMore) ResetWithLock() *sync.Mutex {
    65  	t.mu.Lock()
    66  	defer atomic.StoreUint32(&t.done, 0)
    67  	return &t.mu
    68  }