go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/experiments/cb-feed/processor/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/wcharczuk/go-diskq"
	"github.com/wcharczuk/go-incr"
	"github.com/wcharczuk/go-incr/incrutil"
	"github.com/wcharczuk/go-incr/incrutil/slicei"

	"go.charczuk.com/sdk/ansi"

	"go.charczuk.com/experiments/cb-feed/pkg/coinbase"
	"go.charczuk.com/experiments/cb-feed/pkg/price"
)

var (
	flagPath           = flag.String("path", "_data", "The data stream path")
	flagReplayDuration = flag.Duration("replay", -5*time.Minute, "The duration of time to replay from the data file")
)

func main() {
	flag.Parse()

	partitions, err := diskq.GetPartitions(*flagPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	startOffsets := make(map[uint32]uint64)
	for _, p := range partitions {
		startOffset, _, err := diskq.GetOffsetAfter(*flagPath, p, time.Now().UTC().Add(*flagReplayDuration))
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			os.Exit(1)
		}
		fmt.Printf("starting consumer for partition %d at %d\n", p, startOffset)
		startOffsets[p] = startOffset
	}
	dataReader, err := diskq.OpenConsumerGroup(*flagPath, diskq.ConsumerGroupOptions{
		OnCreateConsumer: func(partitionIndex uint32) (diskq.ConsumerOptions, error) {
			return diskq.ConsumerOptions{
				StartBehavior: diskq.ConsumerStartBehaviorAtOffset,
				StartOffset:   startOffsets[partitionIndex],
				EndBehavior:   diskq.ConsumerEndBehaviorWait,
			}, nil
		},
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	defer dataReader.Close()

	g := incr.New()

	products := make(map[string]*productGraph)
	for _, productID := range [...]string{
		"BTC-USD",
		"ETH-USD",
	} {
		products[productID] = createProductGraph(g, productID)
	}

	ctx, cancel := context.WithCancel(context.Background())

	done := make(chan struct{})
	started := make(chan struct{})
	go func() {
		defer close(done)
		close(started)
		var b coinbase.BasicMessage
		for {
			select {
			case <-ctx.Done():
				return
			case msg, ok := <-dataReader.Messages():
				if !ok {
					return
				}
				_ = json.Unmarshal(msg.Data, &b)
				switch b.Type {
				case "ticker":
					var t coinbase.Ticker
					_ = json.Unmarshal(msg.Data, &t)
					products[t.ProductID].latest.Set(t)
					g.Stabilize(ctx)
				case "match":
					var m coinbase.Match
					_ = json.Unmarshal(msg.Data, &m)
					// g.Stabilize(ctx)
				}
			}
		}
	}()

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)
	<-interrupt
	cancel()
	<-done
}

type productGraph struct {
	productID    string
	latest       incr.VarIncr[coinbase.Ticker]
	extractPrice incr.Incr[price.Price]
	prices       incr.Incr[[]price.Price]
	sortedPrices incr.Incr[[]price.Price]
	latestPrice  incr.Incr[price.Price]
	minPrice     incr.Incr[price.Price]
	maxPrice     incr.Incr[price.Price]

	latestPriceChangeDir incr.Incr[bool]
	minPriceChangeDir    incr.Incr[bool]
	maxPriceChangeDir    incr.Incr[bool]

	observeLatestPriceChangeDir incr.ObserveIncr[bool]
	observeMinPriceChangeDir    incr.ObserveIncr[bool]
	observeMaxPriceChangeDir    incr.ObserveIncr[bool]

	observeLatestPrice incr.ObserveIncr[price.Price]
	observeMinPrice    incr.ObserveIncr[price.Price]
	observeMaxPrice    incr.ObserveIncr[price.Price]
}

func (pg *productGraph) announce(which string) {
	var tokens []string
	if which == "min" {
		if pg.observeMinPriceChangeDir.Value() {
			tokens = append(tokens, ansi.Green(fmt.Sprintf("min=%0.2f", pg.observeMinPrice.Value().Float64())))
		} else {
			tokens = append(tokens, ansi.Red(fmt.Sprintf("min=%0.2f", pg.observeMinPrice.Value().Float64())))
		}
	} else {
		tokens = append(tokens, fmt.Sprintf("min=%0.2f", pg.observeMinPrice.Value().Float64()))
	}
	if which == "max" {
		if pg.observeMaxPriceChangeDir.Value() {
			tokens = append(tokens, ansi.Green(fmt.Sprintf("max=%0.2f", pg.observeMaxPrice.Value().Float64())))
		} else {
			tokens = append(tokens, ansi.Red(fmt.Sprintf("max=%0.2f", pg.observeMaxPrice.Value().Float64())))
		}
	} else {
		tokens = append(tokens, fmt.Sprintf("max=%0.2f", pg.observeMaxPrice.Value().Float64()))
	}
	if which == "latest" {
		if pg.observeLatestPriceChangeDir.Value() {
			tokens = append(tokens, ansi.Green(fmt.Sprintf("latest=%0.2f", pg.observeLatestPrice.Value().Float64())))
		} else {
			tokens = append(tokens, ansi.Red(fmt.Sprintf("latest=%0.2f", pg.observeLatestPrice.Value().Float64())))
		}
	} else {
		tokens = append(tokens, fmt.Sprintf("latest=%0.2f", pg.observeLatestPrice.Value().Float64()))
	}
	fmt.Printf("%s | %s\n",
		pg.productID,
		strings.Join(tokens, " "),
	)
}

func createProductGraph(g *incr.Graph, productID string) *productGraph {
	output := productGraph{
		productID: productID,
	}
	output.latest = incr.Var(g, coinbase.Ticker{})
	output.extractPrice = incr.Map(g, output.latest, func(m coinbase.Ticker) price.Price {
		return price.Parse(m.Price)
	})
	output.prices = slicei.Accumulate(g, output.extractPrice, func(values []price.Price, newValue price.Price) []price.Price {
		return append(values, newValue)
	})
	output.sortedPrices = slicei.AccumulateSorted(g, output.extractPrice, Asc)
	output.latestPrice = slicei.Last(g, output.prices)
	output.minPrice = slicei.First(g, output.sortedPrices)
	output.maxPrice = slicei.Last(g, output.sortedPrices)

	output.latestPriceChangeDir = changeDirection(g, output.latestPrice)
	output.minPriceChangeDir = changeDirection(g, output.minPrice)
	output.maxPriceChangeDir = changeDirection(g, output.maxPrice)

	output.observeLatestPriceChangeDir = incr.MustObserve(g, output.latestPriceChangeDir)
	output.observeLatestPrice = incr.MustObserve(g, incrutil.CutoffUnchanged(g, output.latestPrice))
	output.observeLatestPrice.OnUpdate(func(_ context.Context, p price.Price) {
		output.announce("latest")
	})
	output.observeMinPriceChangeDir = incr.MustObserve(g, output.minPriceChangeDir)
	output.observeMinPrice = incr.MustObserve(g, incrutil.CutoffUnchanged(g, output.minPrice))
	output.observeMinPrice.OnUpdate(func(_ context.Context, p price.Price) {
		output.announce("min")
	})
	output.observeMaxPriceChangeDir = incr.MustObserve(g, output.maxPriceChangeDir)
	output.observeMaxPrice = incr.MustObserve(g, incrutil.CutoffUnchanged(g, output.maxPrice))
	output.observeMaxPrice.OnUpdate(func(_ context.Context, p price.Price) {
		output.announce("max")
	})
	return &output
}

func changeDirection(g *incr.Graph, from incr.Incr[price.Price]) incr.Incr[bool] {
	return incrutil.MapLast(g, from, func(v0, v1 price.Price) bool {
		return v0 < v1
	})
}

type numbers interface {
	~int64 | ~uint64 | ~float64 | ~string
}

// Asc returns a sorted comparer for sortable values in ascending order.
func Asc[A numbers](testValue, newValue A) int {
	if testValue == newValue {
		return 0
	}
	if testValue < newValue {
		return -1
	}
	return 1
}