mirror of
https://github.com/NotAShelf/tct.git
synced 2025-10-02 15:03:48 +00:00
treewide: refactor; defer cli handling to cobra
This commit is contained in:
parent
7a5207f3b4
commit
1dd971e16b
7 changed files with 185 additions and 151 deletions
91
cmd/root.go
Normal file
91
cmd/root.go
Normal file
|
@ -0,0 +1,91 @@
|
|||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"notashelf.dev/tct/internal/tct"
|
||||
)
|
||||
|
||||
var (
|
||||
url string
|
||||
maxRequests int
|
||||
delay time.Duration
|
||||
Version string // will be set by main package
|
||||
)
|
||||
|
||||
var rootCmd = &cobra.Command{
|
||||
Use: "tct",
|
||||
Short: "TCP Connection Timer - find optimal parallel request count",
|
||||
Long: `A tool to measure and find the optimal number of parallel TCP requests for a given URL.`,
|
||||
Run: run,
|
||||
}
|
||||
|
||||
func Execute() {
|
||||
if Version != "" {
|
||||
rootCmd.Version = Version
|
||||
}
|
||||
if err := rootCmd.Execute(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func init() {
|
||||
rootCmd.Flags().StringVarP(&url, "url", "u", "http://example.com", "URL to fetch")
|
||||
rootCmd.Flags().IntVarP(&maxRequests, "max", "m", 100, "Maximum number of parallel requests")
|
||||
rootCmd.Flags().DurationVarP(&delay, "delay", "d", 0, "Delay between requests")
|
||||
}
|
||||
|
||||
func run(cmd *cobra.Command, args []string) {
|
||||
client := tct.NewClient()
|
||||
|
||||
// Context that listens for the interrupt signal
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
optimal := 0
|
||||
minTime := time.Duration(0)
|
||||
|
||||
for i := 1; i <= maxRequests; i++ {
|
||||
// Check for cancellation before starting the loop
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Add(i)
|
||||
|
||||
for range i {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
fmt.Println("Operation cancelled.")
|
||||
return
|
||||
default:
|
||||
time.Sleep(delay)
|
||||
client.MakeRequest(url)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
duration := time.Since(start)
|
||||
fmt.Printf("Parallel Requests: %d, Time Taken: %s\n", i, duration)
|
||||
|
||||
if minTime == 0 || duration < minTime {
|
||||
minTime = duration
|
||||
optimal = i
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("\nOptimal Number of Parallel TCP Requests: %d\n", optimal)
|
||||
}
|
10
go.mod
10
go.mod
|
@ -1,12 +1,10 @@
|
|||
module notashelf.dev/tct
|
||||
|
||||
go 1.22.2
|
||||
go 1.24.3
|
||||
|
||||
require github.com/stretchr/testify v1.9.0
|
||||
require github.com/spf13/cobra v1.9.1
|
||||
|
||||
require (
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/stretchr/objx v0.5.2 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/spf13/pflag v1.0.6 // indirect
|
||||
)
|
||||
|
|
18
go.sum
18
go.sum
|
@ -1,12 +1,10 @@
|
|||
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/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
|
||||
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||
github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
|
||||
github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
|
||||
github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
|
||||
github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
|
|
32
internal/tct/client.go
Normal file
32
internal/tct/client.go
Normal file
|
@ -0,0 +1,32 @@
|
|||
package tct
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type HttpClient interface {
|
||||
Get(url string) (*http.Response, error)
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
httpClient HttpClient
|
||||
}
|
||||
|
||||
func NewClient() *Client {
|
||||
return &Client{
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) MakeRequest(url string) {
|
||||
resp, err := c.httpClient.Get(url)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
_ = body
|
||||
}
|
46
internal/tct/client_test.go
Normal file
46
internal/tct/client_test.go
Normal file
|
@ -0,0 +1,46 @@
|
|||
package tct
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type MockHttpClient struct {
|
||||
GetFunc func(url string) (*http.Response, error)
|
||||
}
|
||||
|
||||
func (m *MockHttpClient) Get(url string) (*http.Response, error) {
|
||||
return m.GetFunc(url)
|
||||
}
|
||||
|
||||
func TestClientMakeRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
}{
|
||||
{
|
||||
name: "valid URL",
|
||||
url: "http://example.com",
|
||||
},
|
||||
{
|
||||
name: "another valid URL",
|
||||
url: "https://google.com",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockClient := &MockHttpClient{
|
||||
GetFunc: func(url string) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: http.NoBody,
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
client := &Client{httpClient: mockClient}
|
||||
client.MakeRequest(tt.url)
|
||||
})
|
||||
}
|
||||
}
|
77
main.go
77
main.go
|
@ -1,79 +1,10 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
import "notashelf.dev/tct/cmd"
|
||||
|
||||
type HttpClient interface {
|
||||
Get(url string) (*http.Response, error)
|
||||
}
|
||||
|
||||
var httpClient HttpClient = &http.Client{}
|
||||
var version = "dev" // will be set by build process
|
||||
|
||||
func main() {
|
||||
urlPtr := flag.String("url", "http://example.com", "URL to fetch")
|
||||
maxRequestsPtr := flag.Int("max", 100, "Maximum number of parallel requests")
|
||||
delayPtr := flag.Duration("delay", 0, "Delay between requests")
|
||||
|
||||
flag.Parse()
|
||||
url := *urlPtr
|
||||
maxRequests := *maxRequestsPtr
|
||||
delay := *delayPtr
|
||||
|
||||
// context that listens for the interrupt signal.
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
optimal := 0
|
||||
minTime := time.Duration(0)
|
||||
for i := 1; i <= maxRequests; i++ {
|
||||
// check for cancellation before starting the loop
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
start := time.Now()
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Add(i)
|
||||
for j := 0; j < i; j++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
fmt.Println("Operation cancelled.")
|
||||
return
|
||||
default:
|
||||
time.Sleep(delay)
|
||||
makeRequest(httpClient, url)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
duration := time.Since(start)
|
||||
fmt.Printf("Parallel Requests: %d, Time Taken: %s\n", i, duration)
|
||||
if minTime == 0 || duration < minTime {
|
||||
minTime = duration
|
||||
optimal = i
|
||||
}
|
||||
}
|
||||
fmt.Printf("\nOptimal Number of Parallel TCP Requests: %d\n", optimal)
|
||||
}
|
||||
|
||||
func makeRequest(client HttpClient, url string) {
|
||||
resp, err := client.Get(url)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
_ = body
|
||||
cmd.Version = version
|
||||
cmd.Execute()
|
||||
}
|
||||
|
|
62
main_test.go
62
main_test.go
|
@ -1,62 +0,0 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
type MockHttpClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockHttpClient) Get(url string) (*http.Response, error) {
|
||||
args := m.Called(url)
|
||||
return args.Get(0).(*http.Response), args.Error(1)
|
||||
}
|
||||
|
||||
func TestMakeRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
responseBody string
|
||||
expectedError error
|
||||
}{
|
||||
{
|
||||
name: "Successful request",
|
||||
statusCode: http.StatusOK,
|
||||
responseBody: "OK",
|
||||
},
|
||||
{
|
||||
name: "Failed request",
|
||||
statusCode: http.StatusInternalServerError,
|
||||
responseBody: "",
|
||||
expectedError: errors.New("request failed"),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockClient := new(MockHttpClient)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(tt.statusCode)
|
||||
w.Write([]byte(tt.responseBody))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
mockClient.On("Get", server.URL).Return(&http.Response{
|
||||
StatusCode: tt.statusCode,
|
||||
Body: io.NopCloser(bytes.NewBufferString(tt.responseBody)),
|
||||
}, tt.expectedError)
|
||||
|
||||
makeRequest(mockClient, server.URL)
|
||||
|
||||
mockClient.AssertExpectations(t)
|
||||
})
|
||||
}
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue