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

     1  // Copyright 2020 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 notify
    15  
    16  import (
    17  	"sync"
    18  	"time"
    19  
    20  	"github.com/pingcap/tiflow/pkg/errors"
    21  )
    22  
    23  // Notifier provides a one-to-many notification mechanism
    24  type Notifier struct {
    25  	receivers []struct {
    26  		rec   *Receiver
    27  		index int
    28  	}
    29  	maxIndex int
    30  	mu       sync.RWMutex
    31  	closed   bool
    32  }
    33  
    34  // Notify sends a signal to the Receivers
    35  func (n *Notifier) Notify() {
    36  	n.mu.RLock()
    37  	defer n.mu.RUnlock()
    38  	for _, receiver := range n.receivers {
    39  		receiver.rec.signalNonBlocking()
    40  	}
    41  }
    42  
    43  // NewReceiver creates a receiver
    44  // returns a channel to receive notifications and a function to close this receiver
    45  func (n *Notifier) NewReceiver(tickTime time.Duration) (*Receiver, error) {
    46  	n.mu.Lock()
    47  	defer n.mu.Unlock()
    48  	if n.closed {
    49  		return nil, errors.ErrOperateOnClosedNotifier.GenWithStackByArgs()
    50  	}
    51  	currentIndex := n.maxIndex
    52  	n.maxIndex++
    53  	receiverCh := make(chan struct{}, 1)
    54  	closeCh := make(chan struct{})
    55  	var ticker *time.Ticker
    56  	if tickTime > 0 {
    57  		ticker = time.NewTicker(tickTime)
    58  	}
    59  	rec := &Receiver{
    60  		C: receiverCh,
    61  		c: receiverCh,
    62  		Stop: func() {
    63  			n.remove(currentIndex)
    64  		},
    65  		ticker:  ticker,
    66  		closeCh: closeCh,
    67  	}
    68  	if tickTime > 0 {
    69  		rec.signalTickLoop()
    70  	}
    71  	n.receivers = append(n.receivers, struct {
    72  		rec   *Receiver
    73  		index int
    74  	}{rec: rec, index: currentIndex})
    75  	return rec, nil
    76  }
    77  
    78  func (n *Notifier) remove(index int) {
    79  	n.mu.Lock()
    80  	defer n.mu.Unlock()
    81  	for i, receiver := range n.receivers {
    82  		if receiver.index == index {
    83  			n.receivers = append(n.receivers[:i], n.receivers[i+1:]...)
    84  			close(receiver.rec.closeCh)
    85  			if receiver.rec.ticker != nil {
    86  				receiver.rec.ticker.Stop()
    87  			}
    88  			break
    89  		}
    90  	}
    91  }
    92  
    93  // Close closes the notify and stops all receiver in this notifier
    94  // Note we must `Close` the notifier if we can't ensure each receiver of this
    95  // notifier is called `Stop` in order to prevent goroutine leak.
    96  func (n *Notifier) Close() {
    97  	n.mu.Lock()
    98  	defer n.mu.Unlock()
    99  	for _, receiver := range n.receivers {
   100  		receiver.rec.close()
   101  	}
   102  	n.receivers = nil
   103  	n.closed = true
   104  }
   105  
   106  // Receiver is a receiver of notifier, including the receiver channel and stop receiver function.
   107  type Receiver struct {
   108  	C       <-chan struct{}
   109  	c       chan<- struct{}
   110  	Stop    func()
   111  	ticker  *time.Ticker
   112  	closeCh chan struct{}
   113  }
   114  
   115  // returns true if the receiverCh should be closed
   116  func (r *Receiver) signalNonBlocking() bool {
   117  	select {
   118  	case <-r.closeCh:
   119  		return true
   120  	case r.c <- struct{}{}:
   121  	default:
   122  	}
   123  	return false
   124  }
   125  
   126  func (r *Receiver) signalTickLoop() {
   127  	go func() {
   128  	loop:
   129  		for {
   130  			select {
   131  			case <-r.closeCh:
   132  				break loop
   133  			case <-r.ticker.C:
   134  			}
   135  			exit := r.signalNonBlocking()
   136  			if exit {
   137  				break loop
   138  			}
   139  		}
   140  		close(r.c)
   141  	}()
   142  }
   143  
   144  func (r *Receiver) close() {
   145  	if r.ticker != nil {
   146  		// in this case, r.c could be accessed by signalTickLoop goroutine, hence
   147  		// we should not close it here.
   148  		r.ticker.Stop()
   149  	} else {
   150  		close(r.c)
   151  	}
   152  	close(r.closeCh)
   153  }