axiom-checker/portfolio.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

const SOL_MINT = "So11111111111111111111111111111111111111112"

type ApiKeyWorker struct {
	apiKey          string
	lastRequestTime time.Time
	mutex           sync.Mutex
}

var (
	apiKeyWorkers []*ApiKeyWorker
	workerIndex   int
	workerMutex   sync.Mutex
)

func initApiKeyWorkers(cfg *Config) {
	apiKeyWorkers = make([]*ApiKeyWorker, len(cfg.JupiterApiKeys))
	for i, key := range cfg.JupiterApiKeys {
		apiKeyWorkers[i] = &ApiKeyWorker{
			apiKey:          key,
			lastRequestTime: time.Time{},
		}
	}
	workerIndex = 0
}

func getNextApiKeyWorker() *ApiKeyWorker {
	workerMutex.Lock()
	defer workerMutex.Unlock()

	worker := apiKeyWorkers[workerIndex]
	workerIndex = (workerIndex + 1) % len(apiKeyWorkers)
	return worker
}

func (w *ApiKeyWorker) rateLimit(cfg *Config) {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	elapsed := time.Since(w.lastRequestTime)
	waitTime := time.Duration(cfg.JupiterRateLimitMs)*time.Millisecond - elapsed

	if waitTime > 0 {
		time.Sleep(waitTime)
	}

	w.lastRequestTime = time.Now()
}

type PortfolioValue struct {
	Value      float64
	TokenCount int
}

type Balance struct {
	UiAmount float64 `json:"uiAmount"`
}

type PriceData struct {
	UsdPrice float64 `json:"usdPrice"`
}

func getPortfolioValue(address string, cfg *Config) (*PortfolioValue, error) {
	const MAX_RETRIES = 2

	// Get a worker for this request
	worker := getNextApiKeyWorker()

	for attempt := 0; attempt < MAX_RETRIES; attempt++ {
		result, err := attemptGetPortfolio(address, cfg, worker)
		if err == nil {
			return result, nil
		}
		if attempt == MAX_RETRIES-1 {
			return &PortfolioValue{Value: 0, TokenCount: 0}, nil
		}
	}

	return &PortfolioValue{Value: 0, TokenCount: 0}, nil
}

func attemptGetPortfolio(address string, cfg *Config, worker *ApiKeyWorker) (*PortfolioValue, error) {
	const BATCH_SIZE = 100
	const MAX_TOKENS = 100

	worker.rateLimit(cfg)

	transport, err := getNextProxy()
	if err != nil {
		return nil, fmt.Errorf("failed to get proxy: %w", err)
	}

	client := &http.Client{
		Transport: transport,
		Timeout:   10 * time.Second,
	}

	// Get balances
	req, err := http.NewRequest("GET", fmt.Sprintf("https://api.jup.ag/ultra/v1/balances/%s", address), nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("x-api-key", worker.apiKey)

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("balances request failed: %d", resp.StatusCode)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var balances map[string]Balance
	if err := json.Unmarshal(body, &balances); err != nil {
		return nil, err
	}

	if len(balances) == 0 {
		return &PortfolioValue{Value: 0, TokenCount: 0}, nil
	}

	// Build token list
	var tokens []string
	for addr := range balances {
		if addr != "SOL" {
			tokens = append(tokens, addr)
		}
	}

	// Add SOL mint if SOL balance exists
	if _, hasSol := balances["SOL"]; hasSol {
		tokens = append(tokens, SOL_MINT)
	}

	if len(tokens) > MAX_TOKENS {
		tokens = tokens[:MAX_TOKENS]
	}

	if len(tokens) == 0 {
		return &PortfolioValue{Value: 0, TokenCount: 0}, nil
	}

	// Split into batches
	var batches [][]string
	for i := 0; i < len(tokens); i += BATCH_SIZE {
		end := i + BATCH_SIZE
		if end > len(tokens) {
			end = len(tokens)
		}
		batches = append(batches, tokens[i:end])
	}

	worker.rateLimit(cfg)

	// Fetch prices for all batches
	allPrices := make(map[string]PriceData)
	for _, batch := range batches {
		ids := strings.Join(batch, "%2C")
		url := fmt.Sprintf("https://api.jup.ag/price/v3?ids=%s", ids)

		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			continue
		}

		req.Header.Set("x-api-key", worker.apiKey)

		resp, err := client.Do(req)
		if err != nil {
			continue
		}

		if resp.StatusCode == http.StatusOK {
			body, err := io.ReadAll(resp.Body)
			resp.Body.Close()
			if err == nil {
				var prices map[string]PriceData
				if json.Unmarshal(body, &prices) == nil {
					for k, v := range prices {
						allPrices[k] = v
					}
				}
			}
		} else {
			resp.Body.Close()
		}
	}

	// Calculate total value
	var total float64
	var count int

	for addr, bal := range balances {
		mint := addr
		if addr == "SOL" {
			mint = SOL_MINT
		}

		price := 0.0
		if priceData, ok := allPrices[mint]; ok {
			price = priceData.UsdPrice
		}

		val := bal.UiAmount * price
		if val > 0.01 {
			total += val
			count++
		}
	}

	return &PortfolioValue{Value: total, TokenCount: count}, nil
}