Skip to content

Commit

Permalink
Adding [slightly better] socket support.
Browse files Browse the repository at this point in the history
  • Loading branch information
strategicpause committed Sep 18, 2023
1 parent bce581e commit 8422730
Show file tree
Hide file tree
Showing 8 changed files with 190 additions and 59 deletions.
117 changes: 117 additions & 0 deletions command/socket/leak.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
package socket

import (
"fmt"
"github.com/strategicpause/memory-leak/metrics"
"golang.org/x/sys/unix"
"time"
)

const (
// MaxConnections defines the maximum number of connections that can be accepted by a given socket. This value is
// defined by /proc/sys/net/core/somaxconn.
MaxConnections = 4096
StartPort = 9090
// By default each established socket connection will write 1 KiB to the buffer.
KiB = 1024
)

var (
LocalAddr = [4]byte{127, 0, 0, 1}
)

type Params struct {
NumSockets int64
NetworkAddressDomain int
ConnectionType int
CommunicationProtocol int
PauseTimeInSeconds time.Duration
}

func tcpLeak(params *Params) error {
PrintParams(params)

var stopChan chan bool
var err error

nextPort := StartPort
currentPort := nextPort

for i := int64(0); i < params.NumSockets; i++ {
// Setup a new service every 4096 connections.
if i%MaxConnections == 0 {
if stopChan != nil {
stopChan <- true
}
currentPort = nextPort
stopChan, err = resetServer(params, currentPort)
if err != nil {
return err
}
nextPort = nextPort + 1
}

clientFd, err := unix.Socket(params.NetworkAddressDomain, params.ConnectionType, params.CommunicationProtocol)
if err != nil {
return err
}

if err = unix.Connect(clientFd, &unix.SockaddrInet4{Port: currentPort, Addr: LocalAddr}); err != nil {
return err
}

_, err = unix.Write(clientFd, make([]byte, KiB))

if i%100 == 0 {
metrics.PrintSocketStats()
time.Sleep(params.PauseTimeInSeconds)
}
}

return nil
}

// resetServer will create a new service socket after 4096 connections (which is the maximum backlog value for listen).
func resetServer(params *Params, port int) (chan bool, error) {
stopChan := make(chan bool, 1)
serverFd := Must(unix.Socket(params.NetworkAddressDomain, params.ConnectionType, params.CommunicationProtocol))

serviceAddr := &unix.SockaddrInet4{
Port: port,
Addr: LocalAddr,
}

if err := unix.Bind(serverFd, serviceAddr); err != nil {
return stopChan, err
}

if err := unix.Listen(serverFd, MaxConnections); err != nil {
return stopChan, err
}
fmt.Printf("Service bound to 127.0.0.1:%d.\n", port)

go func() {
for {
select {
case <-stopChan:
unix.Close(serverFd)
return
default:
unix.Accept(serverFd)
}
}
}()

return stopChan, nil
}

func PrintParams(params *Params) {
fmt.Printf("NumSockets: %v\n", params.NumSockets)
}

func Must[T any](obj T, err error) T {
if err != nil {
panic(err)
}
return obj
}
48 changes: 48 additions & 0 deletions command/socket/register.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package socket

import (
"github.com/urfave/cli"
"golang.org/x/sys/unix"
"time"
)

const (
NumSocketsName = "num-sockets"
CommunicationProtocolName = "comm-protocol"
PauseDurationName = "pause"
)

func Register() cli.Command {
return cli.Command{
Name: "socket",
Usage: "Reproduces a socket leak.",
Action: action,
Flags: flags(),
}
}

func flags() []cli.Flag {
return []cli.Flag{
cli.Int64Flag{
Name: NumSocketsName,
Usage: "specify the number of sockets to create.",
Value: 9223372036854775807,
},
cli.DurationFlag{
Name: PauseDurationName,
Usage: "Time between allocations in seconds.",
Value: time.Second,
},
}
}

func action(ctx *cli.Context) error {
params := &Params{
NumSockets: ctx.Int64(NumSocketsName),
NetworkAddressDomain: unix.AF_INET,
ConnectionType: unix.SOCK_STREAM,
CommunicationProtocol: unix.IPPROTO_TCP,
PauseTimeInSeconds: ctx.Duration(PauseDurationName),
}
return tcpLeak(params)
}
33 changes: 0 additions & 33 deletions command/tcp/leak.go

This file was deleted.

24 changes: 0 additions & 24 deletions command/tcp/register.go

This file was deleted.

2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ go 1.19

require (
github.com/inhies/go-bytesize v0.0.0-20220417184213-4913239db9cf
github.com/prometheus/procfs v0.11.1
github.com/urfave/cli v1.22.14
golang.org/x/sys v0.12.0
)

require (
Expand Down
5 changes: 5 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,13 @@ github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
github.com/inhies/go-bytesize v0.0.0-20220417184213-4913239db9cf h1:FtEj8sfIcaaBfAKrE1Cwb61YDtYq9JxChK1c7AKce7s=
github.com/inhies/go-bytesize v0.0.0-20220417184213-4913239db9cf/go.mod h1:yrqSXGoD/4EKfF26AOGzscPOgTTJcyAwM2rpixWT+t4=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/procfs v0.11.1 h1:xRC8Iq1yyca5ypa9n1EZnWZkt7dwcoRPQwX/5gwaUuI=
github.com/prometheus/procfs v0.11.1/go.mod h1:eesXgaPo1q7lBpVMoMy0ZOFTth9hBn4W/y0/p/ScXhY=
github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
Expand All @@ -19,6 +22,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
github.com/urfave/cli v1.22.14 h1:ebbhrRiGK2i4naQJr+1Xj92HXZCrK7MsyTS/ob3HnAk=
github.com/urfave/cli v1.22.14/go.mod h1:X0eDS6pD6Exaclxm99NJ3FiCDRED7vIHpx2mDOHLvkA=
golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
Expand Down
4 changes: 2 additions & 2 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ package main

import (
"github.com/strategicpause/memory-leak/command/memory"
"github.com/strategicpause/memory-leak/command/tcp"
"github.com/strategicpause/memory-leak/command/socket"
"github.com/urfave/cli"
"log"
"os"
Expand All @@ -21,6 +21,6 @@ func main() {
func RegisterCommands() cli.Commands {
return cli.Commands{
memory.Register(),
tcp.Register(),
socket.Register(),
}
}
16 changes: 16 additions & 0 deletions metrics/metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package metrics

import (
"fmt"
"github.com/prometheus/procfs"
"runtime"
)

Expand Down Expand Up @@ -39,3 +40,18 @@ func BToMiB[T uint64 | uint32](b T) uint64 {
func NsToUs(s uint64) uint64 {
return s / 1000
}

func PrintSocketStats() {
fs, _ := procfs.NewDefaultFS()
stats, _ := fs.NetSockstat()
for _, protocol := range stats.Protocols {
if protocol.Protocol == "TCP" {
fmt.Printf("TCP: inuse %d orphan %d tw %d alloc %d mem %d total %d\n", protocol.InUse,
*protocol.Orphan,
*protocol.TW,
*protocol.Alloc,
*protocol.Mem,
*stats.Used)
}
}
}

0 comments on commit 8422730

Please sign in to comment.