github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/ctxmu/ctxmu.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     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  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  // Copyright 2012-2020, Hǎi-Liàng “Hal” Wáng
    15  // Copyright 2022 PingCAP, Inc.
    16  //
    17  // Use of this source code is governed by a BSD-style
    18  // license that can be found in the LICENSE file or at
    19  // https://developers.google.com/open-source/licenses/bsd
    20  //
    21  // original code: https://h12.io/article/go-pattern-context-aware-lock
    22  
    23  package ctxmu
    24  
    25  import "context"
    26  
    27  // CtxMutex implements a context aware lock
    28  type CtxMutex struct {
    29  	ch chan struct{}
    30  }
    31  
    32  // New creates a new CtxMutex
    33  func New() *CtxMutex {
    34  	return &CtxMutex{
    35  		ch: make(chan struct{}, 1),
    36  	}
    37  }
    38  
    39  // Lock acquires a lock, it can be canceled by context
    40  func (mu *CtxMutex) Lock(ctx context.Context) bool {
    41  	select {
    42  	case <-ctx.Done():
    43  		return false
    44  	case mu.ch <- struct{}{}:
    45  		return true
    46  	}
    47  }
    48  
    49  // Unlock releases the acquired lock
    50  func (mu *CtxMutex) Unlock() {
    51  	<-mu.ch
    52  }
    53  
    54  // Locked checks whether the lock is hold
    55  func (mu *CtxMutex) Locked() bool {
    56  	return len(mu.ch) > 0 // locked or not
    57  }