github.com/hernad/nomad@v1.6.112/nomad/stream/subscription.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package stream
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"sync/atomic"
    10  
    11  	"github.com/hernad/nomad/nomad/structs"
    12  )
    13  
    14  const (
    15  	// subscriptionStateOpen is the default state of a subscription. An open
    16  	// subscription may receive new events.
    17  	subscriptionStateOpen uint32 = 0
    18  
    19  	// subscriptionStateClosed indicates that the subscription was closed, possibly
    20  	// as a result of a change to an ACL token, and will not receive new events.
    21  	// The subscriber must issue a new Subscribe request.
    22  	subscriptionStateClosed uint32 = 1
    23  )
    24  
    25  // ErrSubscriptionClosed is a error signalling the subscription has been
    26  // closed. The client should Unsubscribe, then re-Subscribe.
    27  var ErrSubscriptionClosed = errors.New("subscription closed by server, client should resubscribe")
    28  var ErrACLInvalid = errors.New("Provided ACL token is invalid for requested topics")
    29  
    30  type Subscription struct {
    31  	// state must be accessed atomically 0 means open, 1 means closed with reload
    32  	state uint32
    33  
    34  	req *SubscribeRequest
    35  
    36  	// currentItem stores the current buffer item we are on. It
    37  	// is mutated by calls to Next.
    38  	currentItem *bufferItem
    39  
    40  	// forceClosed is closed when forceClose is called. It is used by
    41  	// EventBroker to cancel Next().
    42  	forceClosed chan struct{}
    43  
    44  	// unsub is a function set by EventBroker that is called to free resources
    45  	// when the subscription is no longer needed.
    46  	// It must be safe to call the function from multiple goroutines and the function
    47  	// must be idempotent.
    48  	unsub func()
    49  }
    50  
    51  type SubscribeRequest struct {
    52  	Token     string
    53  	Index     uint64
    54  	Namespace string
    55  
    56  	Topics map[structs.Topic][]string
    57  
    58  	// StartExactlyAtIndex specifies if a subscription needs to
    59  	// start exactly at the requested Index. If set to false,
    60  	// the closest index in the buffer will be returned if there is not
    61  	// an exact match
    62  	StartExactlyAtIndex bool
    63  }
    64  
    65  func newSubscription(req *SubscribeRequest, item *bufferItem, unsub func()) *Subscription {
    66  	return &Subscription{
    67  		forceClosed: make(chan struct{}),
    68  		req:         req,
    69  		currentItem: item,
    70  		unsub:       unsub,
    71  	}
    72  }
    73  
    74  func (s *Subscription) Next(ctx context.Context) (structs.Events, error) {
    75  	if atomic.LoadUint32(&s.state) == subscriptionStateClosed {
    76  		return structs.Events{}, ErrSubscriptionClosed
    77  	}
    78  
    79  	for {
    80  		next, err := s.currentItem.Next(ctx, s.forceClosed)
    81  		switch {
    82  		case err != nil && atomic.LoadUint32(&s.state) == subscriptionStateClosed:
    83  			return structs.Events{}, ErrSubscriptionClosed
    84  		case err != nil:
    85  			return structs.Events{}, err
    86  		}
    87  		s.currentItem = next
    88  
    89  		events := filter(s.req, next.Events.Events)
    90  		if len(events) == 0 {
    91  			continue
    92  		}
    93  		return structs.Events{Index: next.Events.Index, Events: events}, nil
    94  	}
    95  }
    96  
    97  func (s *Subscription) NextNoBlock() ([]structs.Event, error) {
    98  	if atomic.LoadUint32(&s.state) == subscriptionStateClosed {
    99  		return nil, ErrSubscriptionClosed
   100  	}
   101  
   102  	for {
   103  		next := s.currentItem.NextNoBlock()
   104  		if next == nil {
   105  			return nil, nil
   106  		}
   107  		s.currentItem = next
   108  
   109  		events := filter(s.req, next.Events.Events)
   110  		if len(events) == 0 {
   111  			continue
   112  		}
   113  		return events, nil
   114  	}
   115  }
   116  
   117  func (s *Subscription) Unsubscribe() {
   118  	s.unsub()
   119  }
   120  
   121  // filter events to only those that match a subscriptions topic/keys/namespace
   122  func filter(req *SubscribeRequest, events []structs.Event) []structs.Event {
   123  	if len(events) == 0 {
   124  		return nil
   125  	}
   126  
   127  	allTopicKeys := req.Topics[structs.TopicAll]
   128  
   129  	// Return all events if subscribed to all namespaces and all topics
   130  	if req.Namespace == "*" && len(allTopicKeys) == 1 && allTopicKeys[0] == string(structs.TopicAll) {
   131  		return events
   132  	}
   133  
   134  	var result []structs.Event
   135  
   136  	for _, event := range events {
   137  		if req.Namespace != "*" && event.Namespace != "" && event.Namespace != req.Namespace {
   138  			continue
   139  		}
   140  
   141  		// *[*] always matches
   142  		if len(allTopicKeys) == 1 && allTopicKeys[0] == string(structs.TopicAll) {
   143  			result = append(result, event)
   144  			continue
   145  		}
   146  
   147  		keys := allTopicKeys
   148  
   149  		if topicKeys, ok := req.Topics[event.Topic]; ok {
   150  			keys = append(keys, topicKeys...)
   151  		}
   152  
   153  		if len(keys) == 1 && keys[0] == string(structs.TopicAll) {
   154  			result = append(result, event)
   155  			continue
   156  		}
   157  
   158  		for _, key := range keys {
   159  			if eventMatchesKey(event, key) {
   160  				result = append(result, event)
   161  				continue
   162  			}
   163  		}
   164  	}
   165  
   166  	return result
   167  }
   168  
   169  func eventMatchesKey(event structs.Event, key string) bool {
   170  	if event.Key == key {
   171  		return true
   172  	}
   173  
   174  	for _, fk := range event.FilterKeys {
   175  		if fk == key {
   176  			return true
   177  		}
   178  	}
   179  
   180  	return false
   181  }