handler: add PriceEventHandler

Author: Richard Patel

accounts.go

@@ -20,6 +20,7 @@ import (
 
 	bin "github.com/gagliardetto/binary"
 	"github.com/gagliardetto/solana-go"
+	"github.com/shopspring/decimal"
 )
 
 // Magic is the 32-bit number prefixed on each account.
@@ -137,7 +138,7 @@ type Ema struct {
 	Denom int64
 }
 
-// PriceInfo contains a price adn confidence at a specific slot.
+// PriceInfo contains a price and confidence at a specific slot.
 //
 // This struct can represent either a publisher's contribution or the outcome of price aggregation.
 type PriceInfo struct {
@@ -148,6 +149,25 @@ type PriceInfo struct {
 	PubSlot uint64 // valid publishing slot
 }
 
+func (p *PriceInfo) IsZero() bool {
+	return p == nil || *p == PriceInfo{}
+}
+
+// Value returns the parsed price and conf values.
+//
+// If ok is false, the value is invalid.
+func (p *PriceInfo) Value(exponent int32) (price decimal.Decimal, conf decimal.Decimal, ok bool) {
+	price = decimal.New(p.Price, exponent)
+	conf = decimal.New(int64(p.Conf), exponent)
+	ok = p.Status == PriceStatusTrading
+	return
+}
+
+// HasChanged returns whether there was a change between this and another price info.
+func (p *PriceInfo) HasChanged(other *PriceInfo) bool {
+	return (p == nil) != (other == nil) || p.Status != other.Status || p.PubSlot != other.PubSlot
+}
+
 // Price status.
 const (
 	PriceStatusUnknown = uint32(iota)

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gagliardetto/binary v0.6.1
 	github.com/gagliardetto/solana-go v1.3.1-0.20220222155336-dd0af958252d
 	github.com/prometheus/client_golang v1.12.1
+	github.com/shopspring/decimal v1.3.1
 	github.com/stretchr/testify v1.7.0
 	go.uber.org/zap v1.21.0
 )

handler.go

@@ -0,0 +1,208 @@
+//  Copyright 2022 Blockdaemon Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pyth
+
+import (
+	"sync"
+
+	"github.com/gagliardetto/solana-go"
+	"github.com/shopspring/decimal"
+)
+
+// PriceEventHandler provides a callback-style interface to Pyth updates.
+type PriceEventHandler struct {
+	stream *PriceAccountStream
+
+	callbacksLock sync.Mutex // lock over the callbacks map
+	regNonce      uint64
+	callbacks     map[solana.PublicKey]priceCallbacks
+}
+
+// NewPriceEventHandler creates a new event handler over the stream.
+//
+// A stream must not be re-used between event handlers.
+func NewPriceEventHandler(stream *PriceAccountStream) *PriceEventHandler {
+	handler := &PriceEventHandler{
+		stream:    stream,
+		callbacks: make(map[solana.PublicKey]priceCallbacks),
+	}
+	go handler.consume(stream.Updates())
+	return handler
+}
+
+// Err returns the reason why the underlying price account stream is closed.
+//
+// Will block until the stream has actually closed.
+// Returns nil if closure was expected.
+//
+// After this function returns the event handler will not send any more callbacks.
+// You could use this function as a barrier for any cleanup tasks relating to callbacks.
+func (p *PriceEventHandler) Err() error {
+	return p.stream.Err()
+}
+
+// OnPriceChange registers a callback function to be called
+// whenever the aggregate price of the provided price account changes.
+func (p *PriceEventHandler) OnPriceChange(priceKey solana.PublicKey, callback func(PriceUpdate)) CallbackHandle {
+	p.callbacksLock.Lock()
+	defer p.callbacksLock.Unlock()
+	return p.getPriceCallbacks(priceKey).onPrice.register(p, callback)
+}
+
+// OnComponentChange registers a callback function to be called
+// whenever the price component of the given (price account, publisher account) pair changes.
+func (p *PriceEventHandler) OnComponentChange(priceKey solana.PublicKey, publisher solana.PublicKey, callback func(PriceUpdate)) CallbackHandle {
+	p.callbacksLock.Lock()
+	defer p.callbacksLock.Unlock()
+	return p.getComponentCallbacks(priceKey, publisher).register(p, callback)
+}
+
+func (p *PriceEventHandler) getPriceCallbacks(priceKey solana.PublicKey) priceCallbacks {
+	// requires lock
+	res, ok := p.callbacks[priceKey]
+	if !ok {
+		res.init()
+		p.callbacks[priceKey] = res
+	}
+	return res
+}
+
+func (p *PriceEventHandler) getComponentCallbacks(priceKey solana.PublicKey, publisherKey solana.PublicKey) callbackMap {
+	// requires lock
+	price := p.getPriceCallbacks(priceKey)
+	res, ok := price.componentCallbacks[publisherKey]
+	if !ok {
+		res = make(callbackMap)
+		price.componentCallbacks[publisherKey] = res
+	}
+	return res
+}
+
+func (p *PriceEventHandler) consume(updates <-chan PriceAccountUpdate) {
+	for update := range updates {
+		p.processUpdate(update.Pubkey, update.Price)
+	}
+}
+
+func (p *PriceEventHandler) processUpdate(priceKey solana.PublicKey, acc *PriceAccount) {
+	p.callbacksLock.Lock()
+	defer p.callbacksLock.Unlock()
+
+	callbacks := p.callbacks[priceKey]
+	for _, onPrice := range callbacks.onPrice {
+		onPrice.inform(acc, &acc.Agg)
+	}
+	for _, comp := range acc.Components {
+		if comp.Publisher.IsZero() {
+			continue
+		}
+		compCbs := callbacks.componentCallbacks[comp.Publisher]
+		for _, onPrice := range compCbs {
+			onPrice.inform(acc, &comp.Latest)
+		}
+	}
+}
+
+type priceCallbacks struct {
+	onPrice            callbackMap
+	componentCallbacks map[solana.PublicKey]callbackMap
+}
+
+func (p *priceCallbacks) init() {
+	p.onPrice = make(callbackMap)
+	p.componentCallbacks = make(map[solana.PublicKey]callbackMap)
+}
+
+type callbackMap map[uint64]*callbackRegistration
+
+func (container callbackMap) register(p *PriceEventHandler, callback func(PriceUpdate)) CallbackHandle {
+	// requires lock
+	p.regNonce += 1
+	key := p.regNonce
+
+	handle := CallbackHandle{
+		handler:   p,
+		container: container,
+		key:       key,
+	}
+	container[key] = &callbackRegistration{
+		handle:   handle,
+		callback: callback,
+	}
+	return handle
+}
+
+type callbackRegistration struct {
+	previousInfo *PriceInfo
+	callback     func(PriceUpdate)
+	handle       CallbackHandle
+}
+
+func (r *callbackRegistration) inform(acc *PriceAccount, newInfo *PriceInfo) {
+	if r.previousInfo.HasChanged(newInfo) {
+		r.callback(PriceUpdate{
+			Account:      acc,
+			PreviousInfo: r.previousInfo,
+			CurrentInfo:  newInfo,
+		})
+	}
+	r.previousInfo = newInfo
+}
+
+// PriceUpdate is returned to callbacks when an aggregate or component price has been updated.
+type PriceUpdate struct {
+	Account      *PriceAccount
+	PreviousInfo *PriceInfo
+	CurrentInfo  *PriceInfo
+}
+
+// Previous returns the value of the previously seen price update.
+//
+// If ok is false, the value is invalid.
+func (p PriceUpdate) Previous() (price decimal.Decimal, conf decimal.Decimal, ok bool) {
+	if !p.PreviousInfo.IsZero() && p.Account != nil {
+		p.PreviousInfo.Value(p.Account.Exponent)
+	}
+	return
+}
+
+// Current returns the value of the last price update.
+//
+// If ok is false, the value is invalid.
+func (p PriceUpdate) Current() (price decimal.Decimal, conf decimal.Decimal, ok bool) {
+	if !p.CurrentInfo.IsZero() && p.Account != nil {
+		return p.CurrentInfo.Value(p.Account.Exponent)
+	}
+	return
+}
+
+// CallbackHandle tracks the lifetime of a callback registration.
+type CallbackHandle struct {
+	handler   *PriceEventHandler
+	container callbackMap
+	key       uint64
+}
+
+// Unsubscribe de-registers a callback from the handler.
+//
+// Calling Unsubscribe is optional.
+// The handler calls it automatically when the underlying stream closes.
+func (c CallbackHandle) Unsubscribe() {
+	lock := &c.handler.callbacksLock
+	lock.Lock()
+	defer lock.Unlock()
+
+	delete(c.container, c.key)
+}

handler_test.go

@@ -0,0 +1,44 @@
+//  Copyright 2022 Blockdaemon Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pyth
+
+import (
+	"log"
+	"time"
+
+	"github.com/gagliardetto/solana-go"
+)
+
+func ExamplePriceEventHandler() {
+	// Connect to Pyth on Solana devnet.
+	client := NewClient(Devnet, testRPC, testWS)
+
+	// Open new event stream.
+	stream := client.StreamPriceAccounts()
+	handler := NewPriceEventHandler(stream)
+
+	// Subscribe to price account changes.
+	priceKey := solana.MustPublicKeyFromBase58("J83w4HKfqxwcq3BEMMkPFSppX3gqekLyLJBexebFVkix")
+	handler.OnPriceChange(priceKey, func(info PriceUpdate) {
+		price, conf, ok := info.Current()
+		if ok {
+			log.Printf("Price change: $%s ± $%s", price, conf)
+		}
+	})
+
+	// Close stream after a while.
+	<-time.After(10 * time.Second)
+	stream.Close()
+}

stream.go

@@ -28,6 +28,8 @@ import (
 )
 
 // StreamPriceAccounts creates a new stream of price account updates.
+//
+// It will reconnect automatically if the WebSocket connection breaks or stalls.
 func (c *Client) StreamPriceAccounts() *PriceAccountStream {
 	ctx, cancel := context.WithCancel(context.Background())
 	stream := &PriceAccountStream{
@@ -42,8 +44,9 @@ func (c *Client) StreamPriceAccounts() *PriceAccountStream {
 
 // PriceAccountUpdate is a real-time update carrying a price account change.
 type PriceAccountUpdate struct {
-	Slot uint64
-	*PriceAccount
+	Slot   uint64
+	Pubkey solana.PublicKey
+	Price  *PriceAccount
 }
 
 // PriceAccountStream is an ongoing stream of on-chain price account updates.
@@ -175,8 +178,9 @@ func (p *PriceAccountStream) readNextUpdate(ctx context.Context, sub *ws.Program
 
 	// Send update to channel.
 	msg := PriceAccountUpdate{
-		Slot:         update.Context.Slot,
-		PriceAccount: priceAcc,
+		Slot:   update.Context.Slot,
+		Pubkey: update.Value.Pubkey,
+		Price:  priceAcc,
 	}
 	select {
 	case <-ctx.Done():

stream_test.go

@@ -29,6 +29,6 @@ func ExampleClient_StreamPriceAccounts() {
 	}()
 	// Print updates.
 	for update := range stream.Updates() {
-		fmt.Println(update.Agg.Price)
+		fmt.Println(update.Price.Agg.Price)
 	}
 }