github.com/polygon-io/client-go@v1.16.4/websocket/subscription.go (about)

     1  package polygonws
     2  
     3  import (
     4  	"encoding/json"
     5  	"strings"
     6  
     7  	"github.com/polygon-io/client-go/websocket/models"
     8  	"golang.org/x/exp/maps"
     9  	"golang.org/x/exp/slices"
    10  )
    11  
    12  // subscriptions stores topic subscriptions for resubscribing after disconnect.
    13  type subscriptions map[Topic]set
    14  type set map[string]struct{}
    15  
    16  // add inserts a set of tickers to the given topic.
    17  func (subs subscriptions) add(topic Topic, tickers ...string) {
    18  	_, prefixExists := subs[topic]
    19  	if !prefixExists || slices.Contains(tickers, "*") {
    20  		subs[topic] = make(set)
    21  	}
    22  	for _, t := range tickers {
    23  		subs[topic][t] = struct{}{}
    24  	}
    25  }
    26  
    27  // get retrieves a list of subscription messages based on what has been cached.
    28  func (subs subscriptions) get() []json.RawMessage {
    29  	var msgs []json.RawMessage
    30  	for topic, tickers := range subs {
    31  		msg, err := getSub(models.Subscribe, topic, maps.Keys(tickers)...)
    32  		if err != nil {
    33  			continue // skip malformed messages
    34  		}
    35  		msgs = append(msgs, msg)
    36  	}
    37  	return msgs
    38  }
    39  
    40  // delete removes a set of tickers from the given topic.
    41  func (subs subscriptions) delete(topic Topic, tickers ...string) {
    42  	for _, t := range tickers {
    43  		delete(subs[topic], t)
    44  	}
    45  	if len(subs[topic]) == 0 {
    46  		delete(subs, topic)
    47  	}
    48  }
    49  
    50  // getSub builds a subscription message for a given topic.
    51  func getSub(action models.Action, topic Topic, tickers ...string) (json.RawMessage, error) {
    52  	if len(tickers) == 0 {
    53  		tickers = []string{"*"}
    54  	}
    55  
    56  	var params []string
    57  	for _, ticker := range tickers {
    58  		params = append(params, topic.prefix()+"."+ticker)
    59  	}
    60  
    61  	msg, err := json.Marshal(&models.ControlMessage{
    62  		Action: action,
    63  		Params: strings.Join(params, ","),
    64  	})
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	return msg, nil
    70  }