github.com/anycable/anycable-go@v1.5.1/node/subscription_state.go (about)

     1  package node
     2  
     3  import "sync"
     4  
     5  type SubscriptionState struct {
     6  	channels map[string]map[string]struct{}
     7  	mu       sync.RWMutex
     8  }
     9  
    10  func NewSubscriptionState() *SubscriptionState {
    11  	return &SubscriptionState{channels: make(map[string]map[string]struct{})}
    12  }
    13  
    14  func (st *SubscriptionState) HasChannel(id string) bool {
    15  	st.mu.RLock()
    16  	defer st.mu.RUnlock()
    17  
    18  	_, ok := st.channels[id]
    19  	return ok
    20  }
    21  
    22  func (st *SubscriptionState) AddChannel(id string) {
    23  	st.mu.Lock()
    24  	defer st.mu.Unlock()
    25  
    26  	st.channels[id] = make(map[string]struct{})
    27  }
    28  
    29  func (st *SubscriptionState) RemoveChannel(id string) {
    30  	st.mu.Lock()
    31  	defer st.mu.Unlock()
    32  
    33  	delete(st.channels, id)
    34  }
    35  
    36  func (st *SubscriptionState) Channels() []string {
    37  	st.mu.RLock()
    38  	defer st.mu.RUnlock()
    39  
    40  	keys := make([]string, len(st.channels))
    41  	i := 0
    42  
    43  	for k := range st.channels {
    44  		keys[i] = k
    45  		i++
    46  	}
    47  	return keys
    48  }
    49  
    50  func (st *SubscriptionState) ToMap() map[string][]string {
    51  	st.mu.RLock()
    52  	defer st.mu.RUnlock()
    53  
    54  	res := make(map[string][]string, len(st.channels))
    55  
    56  	for k, v := range st.channels {
    57  		streams := make([]string, len(v))
    58  
    59  		i := 0
    60  		for name := range v {
    61  			streams[i] = name
    62  			i++
    63  		}
    64  
    65  		res[k] = streams
    66  	}
    67  
    68  	return res
    69  }
    70  
    71  func (st *SubscriptionState) AddChannelStream(id string, stream string) {
    72  	st.mu.Lock()
    73  	defer st.mu.Unlock()
    74  
    75  	if _, ok := st.channels[id]; ok {
    76  		st.channels[id][stream] = struct{}{}
    77  	}
    78  }
    79  
    80  func (st *SubscriptionState) RemoveChannelStream(id string, stream string) {
    81  	st.mu.Lock()
    82  	defer st.mu.Unlock()
    83  
    84  	if _, ok := st.channels[id]; ok {
    85  		delete(st.channels[id], stream)
    86  	}
    87  }
    88  
    89  func (st *SubscriptionState) RemoveChannelStreams(id string) []string {
    90  	st.mu.Lock()
    91  	defer st.mu.Unlock()
    92  
    93  	if streamNames, ok := st.channels[id]; ok {
    94  		st.channels[id] = make(map[string]struct{})
    95  
    96  		streams := make([]string, len(streamNames))
    97  
    98  		i := 0
    99  		for key := range streamNames {
   100  			streams[i] = key
   101  			i++
   102  		}
   103  
   104  		return streams
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (st *SubscriptionState) StreamsFor(id string) []string {
   111  	st.mu.RLock()
   112  	defer st.mu.RUnlock()
   113  
   114  	if streamNames, ok := st.channels[id]; ok {
   115  		streams := make([]string, len(streamNames))
   116  
   117  		i := 0
   118  		for key := range streamNames {
   119  			streams[i] = key
   120  			i++
   121  		}
   122  
   123  		return streams
   124  	}
   125  
   126  	return nil
   127  }