github.com/cloudwego/kitex@v0.9.0/pkg/utils/sharedticker.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package utils
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  	"time"
    23  
    24  	"github.com/cloudwego/kitex/pkg/gofunc"
    25  )
    26  
    27  type TickerTask interface {
    28  	Tick()
    29  }
    30  
    31  // NewSharedTicker constructs a SharedTicker with specified interval.
    32  func NewSharedTicker(interval time.Duration) *SharedTicker {
    33  	return &SharedTicker{
    34  		Interval: interval,
    35  		tasks:    map[TickerTask]struct{}{},
    36  		stopChan: make(chan struct{}, 1),
    37  	}
    38  }
    39  
    40  type SharedTicker struct {
    41  	sync.Mutex
    42  	started  bool
    43  	Interval time.Duration
    44  	tasks    map[TickerTask]struct{}
    45  	stopChan chan struct{}
    46  }
    47  
    48  func (t *SharedTicker) Add(b TickerTask) {
    49  	if b == nil {
    50  		return
    51  	}
    52  	t.Lock()
    53  	// Add task
    54  	t.tasks[b] = struct{}{}
    55  	if !t.started {
    56  		t.started = true
    57  		go t.Tick(t.Interval)
    58  	}
    59  	t.Unlock()
    60  }
    61  
    62  func (t *SharedTicker) Delete(b TickerTask) {
    63  	t.Lock()
    64  	// Delete from tasks
    65  	delete(t.tasks, b)
    66  	// no tasks remaining then stop the Tick
    67  	if len(t.tasks) == 0 {
    68  		// unblocked when multi Delete call
    69  		select {
    70  		case t.stopChan <- struct{}{}:
    71  			t.started = false
    72  		default:
    73  		}
    74  	}
    75  	t.Unlock()
    76  }
    77  
    78  func (t *SharedTicker) Closed() bool {
    79  	t.Lock()
    80  	defer t.Unlock()
    81  	return !t.started
    82  }
    83  
    84  func (t *SharedTicker) Tick(interval time.Duration) {
    85  	var wg sync.WaitGroup
    86  	ticker := time.NewTicker(interval)
    87  	defer ticker.Stop()
    88  	for {
    89  		select {
    90  		case <-ticker.C:
    91  			t.Lock()
    92  			for b := range t.tasks {
    93  				wg.Add(1)
    94  				task := b
    95  				gofunc.GoFunc(context.Background(), func() {
    96  					defer wg.Done()
    97  					task.Tick()
    98  				})
    99  			}
   100  			t.Unlock()
   101  			wg.Wait()
   102  		case <-t.stopChan:
   103  			return
   104  		}
   105  	}
   106  }