github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/internal/testEnv/ws-load-tester/loadtester.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"os"
     8  	"sync"
     9  	"sync/atomic"
    10  	"time"
    11  
    12  	"github.com/gobwas/ws"
    13  	"github.com/gobwas/ws/wsutil"
    14  	"github.com/ronaksoft/rony"
    15  	"github.com/ronaksoft/rony/internal/testEnv/pb/service"
    16  	"github.com/ronaksoft/rony/log"
    17  	"github.com/ronaksoft/rony/tools"
    18  	"go.uber.org/zap"
    19  )
    20  
    21  /*
    22     Creation Time: 2019 - Jun - 05
    23     Created by:  (ehsan)
    24     Maintainers:
    25        1.  Ehsan N. Moosa (E2)
    26     Auditor: Ehsan N. Moosa (E2)
    27     Copyright Ronak Software Group 2020
    28  */
    29  
    30  var (
    31  	cntConnected      int32
    32  	cntReceivedPacket int32
    33  	cntWritePacket    int32
    34  	latency           int64
    35  	logger            = log.DefaultLogger
    36  )
    37  
    38  func main() {
    39  	fmt.Println("RUN...")
    40  
    41  	logger.SetLevel(log.InfoLevel)
    42  
    43  	var n, m, port int64
    44  	switch len(os.Args) {
    45  	case 4:
    46  		n = tools.StrToInt64(os.Args[1])
    47  		m = tools.StrToInt64(os.Args[2])
    48  		port = tools.StrToInt64(os.Args[3])
    49  	default:
    50  		fmt.Println("needs 3 args, n m port")
    51  
    52  		return
    53  	}
    54  
    55  	go func() {
    56  		for {
    57  			logger.Info("Stats",
    58  				zap.Int32("Received", cntReceivedPacket),
    59  				zap.Int32("Sent", cntWritePacket),
    60  				zap.Int32("Connected", cntConnected),
    61  			)
    62  			time.Sleep(time.Second)
    63  		}
    64  	}()
    65  
    66  	runClients(int(n), int(m), int(port))
    67  	logger.Info("Final Stats",
    68  		zap.Int32("Received", cntReceivedPacket),
    69  		zap.Int32("Sent", cntWritePacket),
    70  		zap.Int32("Connected", cntConnected),
    71  		zap.Duration("Avg Latency", time.Duration(latency/int64(cntReceivedPacket))),
    72  	)
    73  }
    74  
    75  func runClients(n, m, port int) {
    76  	waitGroup := sync.WaitGroup{}
    77  	for i := 0; i < n; i++ {
    78  		waitGroup.Add(1)
    79  		go func(i int) {
    80  			runClient(&waitGroup, m, port)
    81  			waitGroup.Done()
    82  		}(i)
    83  		time.Sleep(time.Millisecond)
    84  	}
    85  	waitGroup.Wait()
    86  	logger.Debug("Running Clients Finished")
    87  }
    88  
    89  func runClient(wg *sync.WaitGroup, m int, port int) {
    90  	var (
    91  		c   net.Conn
    92  		err error
    93  	)
    94  
    95  	for {
    96  		c, _, _, err = ws.Dialer{}.Dial(context.Background(), fmt.Sprintf("ws://localhost:%d", port))
    97  		if err == nil {
    98  			break
    99  		}
   100  		logger.Warn("error on connect: ", zap.Error(err))
   101  		time.Sleep(time.Millisecond)
   102  	}
   103  	atomic.AddInt32(&cntConnected, 1)
   104  	wg.Add(1)
   105  	go func() {
   106  		time.Sleep(time.Second)
   107  		defer wg.Done()
   108  		var m []wsutil.Message
   109  		for {
   110  			m = m[:0]
   111  			_ = c.SetReadDeadline(time.Now().Add(time.Second * 3))
   112  			startTime := tools.CPUTicks()
   113  			m, err = wsutil.ReadServerMessage(c, m)
   114  			d := time.Duration(tools.CPUTicks() - startTime)
   115  			if err != nil {
   116  				err = c.Close()
   117  				if err != nil {
   118  					logger.Warn("Error on Close", zap.Error(err))
   119  				}
   120  
   121  				return
   122  			}
   123  			atomic.AddInt64(&latency, int64(d))
   124  			atomic.AddInt32(&cntReceivedPacket, int32(len(m)))
   125  		}
   126  	}()
   127  
   128  	echoRequest := service.EchoRequest{
   129  		Int:       21313,
   130  		Timestamp: 42342342342,
   131  	}
   132  	req := rony.PoolMessageEnvelope.Get()
   133  	req.Fill(tools.RandomUint64(0), service.C_SampleEcho, &echoRequest)
   134  	reqBytes, _ := req.Marshal()
   135  	for i := 0; i < m; i++ {
   136  		err := wsutil.WriteClientMessage(c, ws.OpBinary, reqBytes)
   137  		if err != nil {
   138  			logger.Error(err.Error())
   139  			_ = c.Close()
   140  
   141  			return
   142  		} else {
   143  			atomic.AddInt32(&cntWritePacket, 1)
   144  		}
   145  	}
   146  }