AITrading/go-service/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	binance "github.com/adshao/go-binance/v2"
	"github.com/redis/go-redis/v9"
)

type KlineMessage struct {
	Type      string `json:"type"`
	Source    string `json:"source"`
	Symbol    string `json:"symbol"`
	Interval  string `json:"interval"`
	EventTime int64  `json:"event_time"`
	OpenTime  int64  `json:"open_time"`
	CloseTime int64  `json:"close_time"`
	Open      string `json:"open"`
	High      string `json:"high"`
	Low       string `json:"low"`
	Close     string `json:"close"`
	Volume    string `json:"volume"`
	TradeNum  int64  `json:"trade_num"`
	Final     bool   `json:"final"`
}

type RedisPublisher interface {
	Publish(ctx context.Context, channel string, payload []byte) error
	Close() error
}

type NoopPublisher struct{}

func (p *NoopPublisher) Publish(_ context.Context, _ string, _ []byte) error {
	return nil
}

func (p *NoopPublisher) Close() error {
	return nil
}

type RedisChannelPublisher struct {
	client *redis.Client
}

func (p *RedisChannelPublisher) Publish(ctx context.Context, channel string, payload []byte) error {
	return p.client.Publish(ctx, channel, payload).Err()
}

func (p *RedisChannelPublisher) Close() error {
	return p.client.Close()
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	symbol := getEnv("BINANCE_SYMBOL", "btcusdt")
	interval := getEnv("BINANCE_INTERVAL", "1m")
	channel := getEnv("REDIS_CHANNEL", "kline.stream")

	publisher, err := initPublisherFromEnv(ctx)
	if err != nil {
		log.Fatalf("init publisher failed: %v", err)
	}
	defer func() {
		if err := publisher.Close(); err != nil {
			log.Printf("publisher close failed: %v", err)
		}
	}()

	doneC, stopC, err := binance.WsKlineServe(symbol, interval,
		func(event *binance.WsKlineEvent) {
			msg := KlineMessage{
				Type:      "kline",
				Source:    "binance_ws",
				Symbol:    strings.ToLower(event.Symbol),
				Interval:  event.Kline.Interval,
				EventTime: event.Time,
				OpenTime:  event.Kline.StartTime,
				CloseTime: event.Kline.EndTime,
				Open:      event.Kline.Open,
				High:      event.Kline.High,
				Low:       event.Kline.Low,
				Close:     event.Kline.Close,
				Volume:    event.Kline.Volume,
				TradeNum:  int64(event.Kline.TradeNum),
				Final:     event.Kline.IsFinal,
			}

			payload, err := json.Marshal(msg)
			if err != nil {
				log.Printf("marshal failed: %v", err)
				return
			}

			fmt.Println(string(payload))

			if err := publisher.Publish(ctx, channel, payload); err != nil {
				log.Printf("redis publish failed: %v", err)
			}
		},
		func(err error) {
			log.Printf("ws error: %v", err)
		},
	)
	if err != nil {
		log.Fatalf("start ws stream failed: %v", err)
	}

	log.Printf("subscribed binance kline stream: symbol=%s interval=%s channel=%s", symbol, interval, channel)

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	select {
	case <-sig:
		log.Println("received shutdown signal")
		close(stopC)
		cancel()
	case <-doneC:
		log.Println("ws stream closed")
		cancel()
	}
}

func initPublisherFromEnv(ctx context.Context) (RedisPublisher, error) {
	addr := strings.TrimSpace(os.Getenv("REDIS_ADDR"))
	if addr == "" {
		log.Println("REDIS_ADDR not configured, using NoopPublisher")
		return &NoopPublisher{}, nil
	}

	db, err := strconv.Atoi(getEnv("REDIS_DB", "0"))
	if err != nil {
		return nil, fmt.Errorf("invalid REDIS_DB: %w", err)
	}

	client := redis.NewClient(&redis.Options{
		Addr:         addr,
		Password:     os.Getenv("REDIS_PASSWORD"),
		DB:           db,
		DialTimeout:  5 * time.Second,
		ReadTimeout:  3 * time.Second,
		WriteTimeout: 3 * time.Second,
	})

	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("redis ping failed: %w", err)
	}

	log.Printf("connected redis at %s db=%d", addr, db)
	return &RedisChannelPublisher{client: client}, nil
}

func getEnv(key, fallback string) string {
	v := strings.TrimSpace(os.Getenv(key))
	if v == "" {
		return fallback
	}
	return v
}