go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/experiments/cb-feed/data-reader/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"os"
     8  	"os/signal"
     9  	"sync"
    10  	"sync/atomic"
    11  	"time"
    12  
    13  	"github.com/gorilla/websocket"
    14  	"github.com/wcharczuk/go-diskq"
    15  	"go.charczuk.com/experiments/cb-feed/pkg/coinbase"
    16  	"go.charczuk.com/sdk/logutil"
    17  )
    18  
    19  const feedURL = "wss://ws-feed.exchange.coinbase.com"
    20  const dataPath = "_data"
    21  
    22  var messagesReceived, messagesWritten uint64
    23  var last = time.Now()
    24  var lastMessagesReceived, lastMessagesWritten uint64
    25  
    26  func main() {
    27  	l := logutil.New()
    28  
    29  	dq, err := diskq.New(diskq.Config{
    30  		Path:              dataPath,
    31  		PartitionCount:    5,
    32  		RetentionMaxBytes: 1 << 30,
    33  		RetentionMaxAge:   12 * time.Hour,
    34  	})
    35  	maybeFatal(err)
    36  	defer dq.Close()
    37  	go func() {
    38  		for range time.Tick(time.Minute) {
    39  			logutil.Info(l, "vacuuming diskq")
    40  			if err := dq.Vacuum(); err != nil {
    41  				logutil.Error(l, err)
    42  			}
    43  		}
    44  	}()
    45  
    46  	go func() {
    47  		for range time.Tick(5 * time.Second) {
    48  			timeDelta := time.Since(last)
    49  			messagesReceivedDelta := atomic.LoadUint64(&messagesReceived) - lastMessagesReceived
    50  			messagesWrittenDelta := atomic.LoadUint64(&messagesWritten) - lastMessagesWritten
    51  
    52  			timeDeltaSeconds := float64(timeDelta) / float64(time.Second)
    53  			messagesReceivedRate := float64(messagesReceivedDelta) / timeDeltaSeconds
    54  			messagesWrittenRate := float64(messagesWrittenDelta) / timeDeltaSeconds
    55  			logutil.Infof(l, "stats messages received=%0.2f/sec written=%0.2f/sec", messagesReceivedRate, messagesWrittenRate)
    56  
    57  			last = time.Now()
    58  			lastMessagesReceived = atomic.LoadUint64(&messagesReceived)
    59  			lastMessagesWritten = atomic.LoadUint64(&messagesWritten)
    60  		}
    61  	}()
    62  
    63  	interrupt := make(chan os.Signal, 1)
    64  	signal.Notify(interrupt, os.Interrupt)
    65  
    66  	wg := sync.WaitGroup{}
    67  	for _, productID := range products {
    68  		wg.Add(1)
    69  		go readProduct(context.Background(), productID, dq, interrupt, &wg)
    70  	}
    71  	wg.Wait()
    72  }
    73  
    74  var products = [...]string{
    75  	"BTC-USD",
    76  	"ETH-USD",
    77  }
    78  
    79  func readProduct(ctx context.Context, productID string, dq *diskq.Diskq, interrupt <-chan os.Signal, wg *sync.WaitGroup) {
    80  	defer wg.Done()
    81  
    82  	l := logutil.GetLogger(ctx)
    83  	client, _, err := websocket.DefaultDialer.Dial(feedURL, nil)
    84  	maybeFatal(err)
    85  	defer client.Close()
    86  
    87  	logutil.Infof(l, "sending subscribe request")
    88  	err = client.WriteJSON(subscribeRequest{
    89  		Type:       "subscribe",
    90  		ProductIDs: []string{productID},
    91  		Channels: []any{
    92  			"heartbeat",
    93  			"matches",
    94  			channel{Name: "ticker", ProductIDs: []string{productID}},
    95  		},
    96  	})
    97  	maybeFatal(err)
    98  
    99  	readDone := make(chan struct{})
   100  	readStarted := make(chan struct{})
   101  	go func() {
   102  		defer close(readDone)
   103  		close(readStarted)
   104  		for {
   105  			_, messageData, err := client.ReadMessage()
   106  			if err != nil {
   107  				logutil.Error(l, err)
   108  				return
   109  			}
   110  			atomic.AddUint64(&messagesReceived, 1)
   111  			var bm coinbase.BasicMessage
   112  			err = json.Unmarshal(messageData, &bm)
   113  			if err != nil {
   114  				logutil.Error(l, err)
   115  				return
   116  			}
   117  			_, _, err = dq.Push(diskq.Message{
   118  				PartitionKey: bm.ProductID,
   119  				Data:         messageData,
   120  			})
   121  			if err != nil {
   122  				logutil.Error(l, err)
   123  				return
   124  			}
   125  			atomic.AddUint64(&messagesWritten, 1)
   126  		}
   127  	}()
   128  	<-readStarted
   129  	select {
   130  	case <-readDone:
   131  		return
   132  	case <-interrupt:
   133  		err := client.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
   134  		if err != nil {
   135  			fmt.Fprintf(os.Stderr, "%v\n", err)
   136  		}
   137  	}
   138  }
   139  
   140  func maybeFatal(err error) {
   141  	if err != nil {
   142  		fmt.Fprintf(os.Stderr, "%v\n", err)
   143  		os.Exit(1)
   144  	}
   145  }
   146  
   147  type subscribeRequest struct {
   148  	Type       string   `json:"type"`
   149  	ProductIDs []string `json:"product_ids"`
   150  	Channels   []any    `json:"channels"`
   151  }
   152  
   153  type channel struct {
   154  	Name       string   `json:"name"`
   155  	ProductIDs []string `json:"product_ids"`
   156  }