Implement some sort of rate limiting

client/client.go

@@ -1,35 +1,67 @@
 package client
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net"
 	"smpptester/pdu"
+	"smpptester/utils"
 	"time"
+
+	"golang.org/x/time/rate"
 )
 
-const RETRY_TIMEOUT = 1 * time.Second
+var (
+	SMPP_CLIENT_RETRY_TIMEOUT      = 1 * time.Second
+	SMPP_CLIENT_SEND_QUEUE_SIZE    = 1024 * 8
+	SMPP_CLIENT_RECEIVE_QUEUE_SIZE = 1024 * 8
+	SMPP_CLIENT_ASYNC_SEND_WORKERS = 16
+	SMPP_CLIENT_SEND_LIMIT         = 4
+)
 
 type SMPPClient struct {
-	Id        int
-	Connected bool
-	enabled   bool
-	running   bool
-	conn      net.Conn
-	port      string
-	log       log.Logger
-	dcErr     chan (error)
+	Id           int
+	Connected    bool
+	SendQueue    chan pdu.PDU
+	ReceiveQueue chan pdu.PDU
+	SendLimit    *utils.ReactiveValue[int]
+
+	sendLimiter *rate.Limiter
+	enabled     bool
+	running     bool
+	conn        net.Conn
+	port        string
+	log         log.Logger
+	dcErr       chan (error)
 }
 
 func NewSMPPClient(port string, id int) *SMPPClient {
 	client := &SMPPClient{
-		Id:    id,
-		port:  port,
-		log:   log.Logger{},
-		dcErr: make(chan error, 128),
+		Id:           id,
+		SendQueue:    make(chan pdu.PDU, SMPP_CLIENT_SEND_QUEUE_SIZE),
+		ReceiveQueue: make(chan pdu.PDU, SMPP_CLIENT_RECEIVE_QUEUE_SIZE),
+		SendLimit:    utils.NewReactiveValue[int](SMPP_CLIENT_SEND_LIMIT),
+
+		sendLimiter: rate.NewLimiter(rate.Limit(SMPP_CLIENT_SEND_LIMIT), SMPP_CLIENT_SEND_LIMIT),
+		port:        port,
+		log:         log.Logger{},
+		dcErr:       make(chan error, 128),
 	}
-	client.log = *log.New(log.Writer(), "", log.LstdFlags)
+	client.log = *log.New(log.Writer(), "", log.LstdFlags|log.Lshortfile)
 	client.log.SetPrefix(fmt.Sprintf("SMPP client %d: ", client.Id))
+
+	go func() {
+		limit := client.SendLimit.Subscribe()
+		defer client.SendLimit.Unsubscribe(limit)
+		for {
+			v := <-limit
+			client.log.Printf("Send limit changed to %d", *v)
+			client.sendLimiter.SetLimit(rate.Limit(*v))
+			client.sendLimiter.SetBurst(*v)
+		}
+	}()
+
 	return client
 }
 
@@ -78,8 +110,8 @@ func (c *SMPPClient) start() {
 			err := <-c.dcErr
 			c.Connected = false
 			if c.enabled {
-				c.log.Printf("Disconnected: '%v' trying again in %d second(s)", err, int(RETRY_TIMEOUT.Seconds()))
-				time.Sleep(RETRY_TIMEOUT)
+				c.log.Printf("Disconnected: '%v' trying again in %d second(s)", err, int(SMPP_CLIENT_RETRY_TIMEOUT.Seconds()))
+				time.Sleep(SMPP_CLIENT_RETRY_TIMEOUT)
 			} else {
 				c.log.Printf("Disconnected: '%v' & client is disabled, quitting", err)
 				break
@@ -117,12 +149,17 @@ func (c *SMPPClient) Send(pdata pdu.PDU) error {
 		return fmt.Errorf("failed to encode PDU: %w", err)
 	}
 
+	err = c.sendLimiter.Wait(context.Background())
+	if err != nil {
+		return fmt.Errorf("failed to wait for rate limiter: %w", err)
+	}
+
 	_, err = c.conn.Write(buf.Bytes())
 	if err != nil {
 		c.dcErr <- err
 		return fmt.Errorf("failed to send PDU: %w", err)
 	}
 
-	c.log.Printf("SMPP client %d sent PDU: %+v", c.Id, pdata)
+	// c.log.Printf("SMPP client %d sent PDU: %+v", c.Id, pdata)
 	return nil
 }