github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/components/statistics/notify.go (about)

     1  package statistics
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/protos/statistic"
     8  	gs "github.com/Asutorufa/yuhaiin/pkg/protos/statistic/grpc"
     9  	"github.com/Asutorufa/yuhaiin/pkg/utils/id"
    10  	"github.com/Asutorufa/yuhaiin/pkg/utils/slice"
    11  	"github.com/Asutorufa/yuhaiin/pkg/utils/syncmap"
    12  )
    13  
    14  type notify struct {
    15  	mu sync.RWMutex
    16  
    17  	notifierIDSeed id.IDGenerator
    18  	notifier       syncmap.SyncMap[uint64, gs.Connections_NotifyServer]
    19  
    20  	channel chan *gs.NotifyData
    21  	closed  context.Context
    22  	close   context.CancelFunc
    23  }
    24  
    25  func newNotify() *notify {
    26  	ctx, cancel := context.WithCancel(context.Background())
    27  	n := &notify{
    28  		channel: make(chan *gs.NotifyData, 100),
    29  		closed:  ctx,
    30  		close:   cancel,
    31  	}
    32  
    33  	go n.start()
    34  
    35  	return n
    36  }
    37  
    38  func (n *notify) register(s gs.Connections_NotifyServer, conns ...connection) uint64 {
    39  	id := n.notifierIDSeed.Generate()
    40  	n.notifier.Store(id, s)
    41  	_ = s.Send(&gs.NotifyData{
    42  		Data: &gs.NotifyData_NotifyNewConnections{
    43  			NotifyNewConnections: &gs.NotifyNewConnections{
    44  				Connections: slice.To(conns, func(c connection) *statistic.Connection { return c.Info() }),
    45  			},
    46  		},
    47  	})
    48  	return id
    49  }
    50  
    51  func (n *notify) unregister(id uint64) { n.notifier.Delete(id) }
    52  
    53  func (n *notify) start() {
    54  	for {
    55  		select {
    56  		case <-n.closed.Done():
    57  			close(n.channel)
    58  			return
    59  		case d := <-n.channel:
    60  			n.notifier.Range(func(key uint64, value gs.Connections_NotifyServer) bool {
    61  				_ = value.Send(d)
    62  				return true
    63  			})
    64  		}
    65  	}
    66  }
    67  
    68  func (n *notify) pubNewConns(conns ...connection) {
    69  	if len(conns) == 0 {
    70  		return
    71  	}
    72  
    73  	n.mu.RLock()
    74  	defer n.mu.RUnlock()
    75  
    76  	select {
    77  	case <-n.closed.Done():
    78  		return
    79  	default:
    80  	}
    81  
    82  	n.channel <- &gs.NotifyData{
    83  		Data: &gs.NotifyData_NotifyNewConnections{
    84  			NotifyNewConnections: &gs.NotifyNewConnections{
    85  				Connections: slice.To(conns, func(c connection) *statistic.Connection { return c.Info() }),
    86  			},
    87  		},
    88  	}
    89  }
    90  
    91  func (n *notify) pubRemoveConns(ids ...uint64) {
    92  	if len(ids) == 0 {
    93  		return
    94  	}
    95  
    96  	n.mu.RLock()
    97  	defer n.mu.RUnlock()
    98  
    99  	select {
   100  	case <-n.closed.Done():
   101  		return
   102  	default:
   103  	}
   104  
   105  	n.channel <- &gs.NotifyData{
   106  		Data: &gs.NotifyData_NotifyRemoveConnections{
   107  			NotifyRemoveConnections: &gs.NotifyRemoveConnections{
   108  				Ids: ids,
   109  			},
   110  		},
   111  	}
   112  }
   113  
   114  func (n *notify) Close() error {
   115  	n.mu.Lock()
   116  	defer n.mu.Unlock()
   117  
   118  	n.close()
   119  	return nil
   120  }