github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/delay/chan.go (about)

     1  package delay
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  const defaultKey = "_default"
    10  
    11  func NewChan(ctx context.Context, fn func(k, v interface{}), delay time.Duration) *Chan {
    12  	d := &Chan{fn: fn, Map: &sync.Map{}, wg: &sync.WaitGroup{}, stop: make(chan struct{}, 1)}
    13  	d.Map.Store(defaultKey, make(chan interface{}, 1))
    14  	d.wg.Add(1)
    15  	go d.run(ctx, delay)
    16  	return d
    17  }
    18  
    19  type Chan struct {
    20  	fn   func(k, v interface{})
    21  	Map  *sync.Map
    22  	wg   *sync.WaitGroup
    23  	stop chan struct{}
    24  }
    25  
    26  func (c *Chan) run(ctx context.Context, delay time.Duration) {
    27  	defer c.wg.Done()
    28  	ticker := time.NewTicker(delay)
    29  	defer ticker.Stop()
    30  	defer c.consume()
    31  
    32  	for {
    33  		select {
    34  		case <-ticker.C:
    35  			c.consume()
    36  		case <-c.stop:
    37  			return
    38  		case <-ctx.Done():
    39  			return
    40  		}
    41  	}
    42  }
    43  
    44  func (c *Chan) Close() error {
    45  	c.stop <- struct{}{}
    46  	c.wg.Wait()
    47  	return nil
    48  }
    49  
    50  func (c *Chan) consume() {
    51  	c.Map.Range(func(k, value interface{}) bool {
    52  		select {
    53  		case v := <-value.(chan interface{}):
    54  			c.fn(k, v)
    55  		default:
    56  		}
    57  		return true
    58  	})
    59  }
    60  
    61  func (c *Chan) PutKey(k string, v interface{}) {
    62  	if ch, ok := c.Map.Load(k); ok {
    63  		replace(ch.(chan interface{}), v)
    64  		return
    65  	}
    66  
    67  	ch, _ := c.Map.LoadOrStore(k, make(chan interface{}, 1))
    68  	replace(ch.(chan interface{}), v)
    69  }
    70  
    71  func (c *Chan) Put(v interface{}) { c.PutKey(defaultKey, v) }
    72  
    73  func replace(ch chan interface{}, v interface{}) {
    74  	// try to remove old one.
    75  	select {
    76  	case <-ch:
    77  	default:
    78  	}
    79  
    80  	// try to replace the new one.
    81  	select {
    82  	case ch <- v:
    83  	default:
    84  	}
    85  }