Implement some sort of rate limiting
This commit is contained in:
@@ -1,18 +1,33 @@
|
|||||||
package client
|
package client
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log"
|
"log"
|
||||||
"net"
|
"net"
|
||||||
"smpptester/pdu"
|
"smpptester/pdu"
|
||||||
|
"smpptester/utils"
|
||||||
"time"
|
"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 {
|
type SMPPClient struct {
|
||||||
Id int
|
Id int
|
||||||
Connected bool
|
Connected bool
|
||||||
|
SendQueue chan pdu.PDU
|
||||||
|
ReceiveQueue chan pdu.PDU
|
||||||
|
SendLimit *utils.ReactiveValue[int]
|
||||||
|
|
||||||
|
sendLimiter *rate.Limiter
|
||||||
enabled bool
|
enabled bool
|
||||||
running bool
|
running bool
|
||||||
conn net.Conn
|
conn net.Conn
|
||||||
@@ -24,12 +39,29 @@ type SMPPClient struct {
|
|||||||
func NewSMPPClient(port string, id int) *SMPPClient {
|
func NewSMPPClient(port string, id int) *SMPPClient {
|
||||||
client := &SMPPClient{
|
client := &SMPPClient{
|
||||||
Id: id,
|
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,
|
port: port,
|
||||||
log: log.Logger{},
|
log: log.Logger{},
|
||||||
dcErr: make(chan error, 128),
|
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))
|
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
|
return client
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -78,8 +110,8 @@ func (c *SMPPClient) start() {
|
|||||||
err := <-c.dcErr
|
err := <-c.dcErr
|
||||||
c.Connected = false
|
c.Connected = false
|
||||||
if c.enabled {
|
if c.enabled {
|
||||||
c.log.Printf("Disconnected: '%v' trying again in %d second(s)", err, int(RETRY_TIMEOUT.Seconds()))
|
c.log.Printf("Disconnected: '%v' trying again in %d second(s)", err, int(SMPP_CLIENT_RETRY_TIMEOUT.Seconds()))
|
||||||
time.Sleep(RETRY_TIMEOUT)
|
time.Sleep(SMPP_CLIENT_RETRY_TIMEOUT)
|
||||||
} else {
|
} else {
|
||||||
c.log.Printf("Disconnected: '%v' & client is disabled, quitting", err)
|
c.log.Printf("Disconnected: '%v' & client is disabled, quitting", err)
|
||||||
break
|
break
|
||||||
@@ -117,12 +149,17 @@ func (c *SMPPClient) Send(pdata pdu.PDU) error {
|
|||||||
return fmt.Errorf("failed to encode PDU: %w", err)
|
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())
|
_, err = c.conn.Write(buf.Bytes())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.dcErr <- err
|
c.dcErr <- err
|
||||||
return fmt.Errorf("failed to send PDU: %w", 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
|
return nil
|
||||||
}
|
}
|
||||||
|
5
go.mod
5
go.mod
@@ -2,4 +2,7 @@ module smpptester
|
|||||||
|
|
||||||
go 1.22.4
|
go 1.22.4
|
||||||
|
|
||||||
require github.com/yuin/gopher-lua v1.1.1
|
require (
|
||||||
|
github.com/yuin/gopher-lua v1.1.1
|
||||||
|
golang.org/x/time v0.5.0
|
||||||
|
)
|
||||||
|
2
go.sum
2
go.sum
@@ -1,2 +1,4 @@
|
|||||||
github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
|
github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
|
||||||
github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
|
github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
|
||||||
|
golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
|
||||||
|
golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
||||||
|
14
main.go
14
main.go
@@ -175,12 +175,16 @@ func main() {
|
|||||||
client.Enable()
|
client.Enable()
|
||||||
for {
|
for {
|
||||||
if client.Connected {
|
if client.Connected {
|
||||||
err := client.Send(submit)
|
// err := client.Send(submit)
|
||||||
if err != nil {
|
// if err != nil {
|
||||||
log.Printf("Failed to send PDU: %v", err)
|
// log.Printf("Failed to send PDU: %v", err)
|
||||||
|
// }
|
||||||
|
go client.Send(submit)
|
||||||
|
// if rand.Int() % 1000 == 0 {
|
||||||
|
// client.SendLimit.Set(rand.Int() % 1000)
|
||||||
|
// }
|
||||||
}
|
}
|
||||||
}
|
time.Sleep(500 * time.Millisecond)
|
||||||
time.Sleep(10 * time.Millisecond)
|
|
||||||
}
|
}
|
||||||
log.Println("Are we done")
|
log.Println("Are we done")
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user