github.com/bitfinexcom/bitfinex-api-go@v0.0.0-20210608095005-9e0b26f200fb/tests/integration/v2/stress_test/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/common"
    10  	"github.com/bitfinexcom/bitfinex-api-go/v2/websocket"
    11  	"github.com/op/go-logging"
    12  )
    13  
    14  var symbols = []string{
    15  	"tBTCUSD", "tETHUSD", "tEOSUSD", "tVETUSD", "tDGBUSD", "tXRPBTC", "tTRXUSD",
    16  	"tLEOUSD", "tLEOBTC", "tLEOUST", "tBTCUST", "tETHBTC", "tETHUST", "tXRPUSD",
    17  }
    18  
    19  func main() {
    20  	useChannels := flag.Bool("channels", false, "subscribes to a lot of channels")
    21  	useMultiplexor := flag.Bool("multiplexer", false, "subscribes/un-subscribes forcing ws connection re-shuffling")
    22  	runTime := flag.Int64("time", 5, "runtime of the stress test in minutes")
    23  	logMode := flag.String("log-level", "INFO", "level of logging. Can be INFO, DEBUG, WARN or ERROR.")
    24  	flag.Parse()
    25  
    26  	p := websocket.NewDefaultParameters()
    27  	logger := logging.MustGetLogger("bfx-websocket")
    28  	p.ManageOrderbook = true
    29  	p.Logger = logger
    30  	logLevel, err := logging.LogLevel(*logMode)
    31  	if err != nil {
    32  		panic(err)
    33  	}
    34  	logging.SetLevel(logLevel, "bfx-websocket")
    35  	c := websocket.NewWithParams(p)
    36  	err = c.Connect()
    37  	if err != nil {
    38  		panic(err)
    39  	}
    40  
    41  	quit := make(chan interface{})
    42  	if *useMultiplexor {
    43  		go runMultiplexerStressTest(c, logger, quit)
    44  	} else if *useChannels {
    45  		go runChannelsStressTest(c, logger, quit)
    46  	} else {
    47  		fmt.Println("No flags set, use:")
    48  		flag.PrintDefaults()
    49  	}
    50  
    51  	timeout := time.After(time.Minute * time.Duration(*runTime))
    52  	for {
    53  		select {
    54  		case msg := <-c.Listen():
    55  			switch msg.(type) {
    56  			case error:
    57  				logger.Error(msg)
    58  			default:
    59  				logger.Debugf("MSG RECV: %#v", msg)
    60  			}
    61  		case <-timeout:
    62  			logger.Warningf("Test timeout of %d mins reached. Killing test.", *runTime)
    63  			panic("time reached")
    64  		case <-quit:
    65  			return
    66  		}
    67  	}
    68  }
    69  
    70  func runChannelsStressTest(client *websocket.Client, log *logging.Logger, quit chan interface{}) {
    71  	fmt.Println("Starting channel stress test...")
    72  	for _, ticker := range symbols {
    73  		log.Infof("Subscribing to trades (%s) (socketCount=%d)", ticker, client.ConnectionCount())
    74  		_, err := client.SubscribeTrades(context.Background(), ticker)
    75  		if err != nil {
    76  			panic(fmt.Sprintf("could not subscribe to trades: %s", err.Error()))
    77  		}
    78  		_, err = client.SubscribeCandles(context.Background(), ticker, common.FifteenMinutes)
    79  		if err != nil {
    80  			panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.FifteenMinutes))
    81  		}
    82  		_, err = client.SubscribeCandles(context.Background(), ticker, common.ThirtyMinutes)
    83  		if err != nil {
    84  			panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.ThirtyMinutes))
    85  		}
    86  		_, err = client.SubscribeCandles(context.Background(), ticker, common.OneHour)
    87  		if err != nil {
    88  			panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.OneHour))
    89  		}
    90  		_, err = client.SubscribeCandles(context.Background(), ticker, common.OneMinute)
    91  		if err != nil {
    92  			panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.OneMinute))
    93  		}
    94  	}
    95  }
    96  
    97  func runMultiplexerStressTest(client *websocket.Client, log *logging.Logger, quit chan interface{}) {
    98  	fmt.Println("Starting multiplexer stress test...")
    99  	for {
   100  		subIds := make([]string, 0)
   101  		for _, ticker := range symbols {
   102  			log.Infof("Subscribing to trades (%s) (socketCount=%d)", ticker, client.ConnectionCount())
   103  			subId1, err := client.SubscribeTrades(context.Background(), ticker)
   104  			if err != nil {
   105  				panic(fmt.Sprintf("could not subscribe to trades: %s", err.Error()))
   106  			}
   107  			subIds = append(subIds, subId1)
   108  			subId2, err := client.SubscribeCandles(context.Background(), ticker, common.FifteenMinutes)
   109  			if err != nil {
   110  				panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.FifteenMinutes))
   111  			}
   112  			subIds = append(subIds, subId2)
   113  			subId3, err := client.SubscribeBook(context.Background(), ticker, common.Precision0, common.FrequencyRealtime, 25)
   114  			if err != nil {
   115  				panic(fmt.Sprintf("could not subscribe to candles %s: %s", err.Error(), common.FifteenMinutes))
   116  			}
   117  			subIds = append(subIds, subId3)
   118  		}
   119  		// wait for a set amount of time before un-subscribing
   120  		time.Sleep(time.Second * 10)
   121  
   122  		// un-subscribe from all channels
   123  		for _, subId := range subIds {
   124  			log.Infof("Un-subscribing from (%s) (socketCount=%d)", subId, client.ConnectionCount())
   125  			err := client.Unsubscribe(context.Background(), subId)
   126  			if err != nil {
   127  				panic(fmt.Sprintf("Could not un-subscribe from channel %s", subId))
   128  			}
   129  		}
   130  
   131  		// wait again
   132  		time.Sleep(time.Second * 10)
   133  	}
   134  }