github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/errctx/center.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  package errctx
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  
    20  	"github.com/pingcap/log"
    21  	"go.uber.org/zap"
    22  )
    23  
    24  // ErrCenter is used to receive errors and provide
    25  // ways to detect the error(s).
    26  type ErrCenter struct {
    27  	rwm      sync.RWMutex
    28  	firstErr error
    29  	children map[*errCtx]struct{}
    30  }
    31  
    32  // NewErrCenter creates a new ErrCenter.
    33  func NewErrCenter() *ErrCenter {
    34  	return &ErrCenter{}
    35  }
    36  
    37  func (c *ErrCenter) removeChild(child *errCtx) {
    38  	c.rwm.Lock()
    39  	defer c.rwm.Unlock()
    40  
    41  	delete(c.children, child)
    42  }
    43  
    44  // OnError receivers an error, if the error center has received one, drops the
    45  // new error and records a warning log. Otherwise the error will be recorded and
    46  // doneCh will be closed to use as notification.
    47  func (c *ErrCenter) OnError(err error) {
    48  	if err == nil {
    49  		return
    50  	}
    51  
    52  	c.rwm.Lock()
    53  	defer c.rwm.Unlock()
    54  
    55  	if c.firstErr != nil {
    56  		// OnError is no-op after the first call with
    57  		// a non-nil error.
    58  		log.Warn("More than one error is received",
    59  			zap.Error(err))
    60  		return
    61  	}
    62  	c.firstErr = err
    63  	for child := range c.children {
    64  		child.doCancel(c.firstErr)
    65  	}
    66  	c.children = nil
    67  }
    68  
    69  // CheckError retusn the recorded error
    70  func (c *ErrCenter) CheckError() error {
    71  	c.rwm.RLock()
    72  	defer c.rwm.RUnlock()
    73  
    74  	return c.firstErr
    75  }
    76  
    77  // WithCancelOnFirstError creates an error context which will cancel the context when the first error is received.
    78  func (c *ErrCenter) WithCancelOnFirstError(ctx context.Context) (context.Context, context.CancelFunc) {
    79  	ec := newErrCtx(ctx)
    80  
    81  	c.rwm.Lock()
    82  	defer c.rwm.Unlock()
    83  
    84  	if c.firstErr != nil {
    85  		// First error is received, cancel the context directly.
    86  		ec.doCancel(c.firstErr)
    87  	} else {
    88  		if c.children == nil {
    89  			c.children = make(map[*errCtx]struct{})
    90  		}
    91  		c.children[ec] = struct{}{}
    92  	}
    93  	return ec, func() {
    94  		ec.doCancel(context.Canceled)
    95  		c.removeChild(ec)
    96  	}
    97  }