bbgo_origin/examples/kucoin/websocket.go

package main

import (
	"context"
	"errors"
	"os"
	"os/signal"
	"time"

	"github.com/c9s/bbgo/pkg/exchange/kucoin"
	"github.com/c9s/bbgo/pkg/exchange/kucoin/kucoinapi"
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

func init() {
	rootCmd.AddCommand(websocketCmd)
}

var websocketCmd = &cobra.Command{
	Use: "websocket",

	// SilenceUsage is an option to silence usage when an error occurs.
	SilenceUsage: true,

	Args: cobra.ExactArgs(1),

	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return nil
		}

		var ctx = context.Background()
		var t = args[0]
		var err error
		var bullet *kucoinapi.Bullet

		switch t {
		case "public":
			bullet, err = client.BulletService.NewGetPublicBulletRequest().Do(ctx)
			if err != nil {
				return err
			}

			logrus.Infof("public bullet: %+v", bullet)

		case "private":
			bullet, err = client.BulletService.NewGetPrivateBulletRequest().Do(ctx)
			if err != nil {
				return err
			}

			logrus.Infof("private bullet: %+v", bullet)

		default:
			return errors.New("valid bullet type: public, private")

		}

		u, err := bullet.URL()
		if err != nil {
			return err
		}

		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)

		logrus.Infof("connecting %s", u.String())
		c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
		if err != nil {
			return err
		}

		defer c.Close()

		id := time.Now().UnixNano() / int64(time.Millisecond)
		wsCmds := []kucoin.WebSocketCommand{
			/*
				{
					Id:             id+1,
					Type:           "subscribe",
					Topic:          "/market/ticker:ETH-USDT",
					PrivateChannel: false,
					Response:       true,
				},
			*/
			{
				Id:             id + 2,
				Type:           "subscribe",
				Topic:          "/market/candles:ETH-USDT_1min",
				PrivateChannel: false,
				Response:       true,
			},
		}

		for _, wsCmd := range wsCmds {
			err = c.WriteJSON(wsCmd)
			if err != nil {
				return err
			}
		}

		done := make(chan struct{})
		go func() {
			defer close(done)
			for {
				_, message, err := c.ReadMessage()
				if err != nil {
					logrus.Infoln("read:", err)
					return
				}

				logrus.Infof("recv: %s", message)
			}
		}()

		pingTicker := time.NewTicker(bullet.PingInterval())
		defer pingTicker.Stop()

		for {
			select {
			case <-done:
				return nil

			case <-pingTicker.C:
				if err := c.WriteJSON(kucoin.WebSocketCommand{
					Id:   time.Now().UnixNano() / int64(time.Millisecond),
					Type: "ping",
				}); err != nil {
					logrus.WithError(err).Error("websocket ping error", err)
				}

			case <-interrupt:
				logrus.Infof("interrupt")

				// Cleanly close the connection by sending a close message and then
				// waiting (with timeout) for the server to close the connection.
				err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
				if err != nil {
					logrus.Error("write close:", err)
					return nil
				}

				select {
				case <-done:
				case <-time.After(time.Second):
				}
				return nil
			}
		}

		return nil
	},
}