.
This commit is contained in:
180
cmd/gembench/gembench.go
Normal file
180
cmd/gembench/gembench.go
Normal file
@@ -0,0 +1,180 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gemserve/lib/logging"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Parse command-line flags
|
||||
insecure := flag.Bool("insecure", false, "Skip TLS certificate verification")
|
||||
totalConnections := flag.Int("total-connections", 250, "Total connections to make")
|
||||
parallelism := flag.Int("parallelism", 10, "How many connections to run in parallel")
|
||||
flag.Parse()
|
||||
|
||||
// Get the URL from arguments
|
||||
args := flag.Args()
|
||||
if len(args) != 1 {
|
||||
fmt.Fprintf(os.Stderr, "Usage: gemget [--insecure] <gemini-url>\n")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
logging.SetupLogging(slog.LevelInfo)
|
||||
logger := logging.Logger
|
||||
ctx := logging.WithLogger(context.Background(), logger)
|
||||
|
||||
geminiURL := args[0]
|
||||
host := validateUrl(geminiURL)
|
||||
if host == "" {
|
||||
logger.Error("Invalid URL.")
|
||||
os.Exit(1)
|
||||
}
|
||||
start := time.Now()
|
||||
err := benchmark(ctx, geminiURL, host, *insecure, *totalConnections, *parallelism)
|
||||
if err != nil {
|
||||
logger.Error(err.Error())
|
||||
os.Exit(1)
|
||||
}
|
||||
end := time.Now()
|
||||
tookMs := end.Sub(start).Milliseconds()
|
||||
logger.Info("End.", "ms", tookMs)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
type ctxKey int
|
||||
|
||||
const ctxKeyJobIndex ctxKey = 1
|
||||
|
||||
func benchmark(ctx context.Context, u string, h string, insecure bool, totalConnections int, parallelism int) error {
|
||||
logger := logging.FromContext(ctx)
|
||||
signals := make(chan os.Signal, 1)
|
||||
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
|
||||
// Root context, used to cancel
|
||||
// connections and graceful shutdown.
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// Semaphore to limit concurrency.
|
||||
// Goroutines put value to channel (acquire slot)
|
||||
// and consume value from channel (release slot).
|
||||
semaphore := make(chan struct{}, parallelism)
|
||||
|
||||
loop:
|
||||
for i := 0; i < totalConnections; i++ {
|
||||
select {
|
||||
case <-signals:
|
||||
logger.Warn("Received SIGINT or SIGTERM signal, shutting down gracefully")
|
||||
cancel()
|
||||
break loop
|
||||
case semaphore <- struct{}{}: // Acquire slot
|
||||
wg.Add(1)
|
||||
go func(jobIndex int) {
|
||||
defer func() {
|
||||
<-semaphore // Release slot
|
||||
wg.Done()
|
||||
}()
|
||||
ctxWithValue := context.WithValue(ctx, ctxKeyJobIndex, jobIndex)
|
||||
ctxWithTimeout, cancel := context.WithTimeout(ctxWithValue, 60*time.Second)
|
||||
defer cancel()
|
||||
err := connect(ctxWithTimeout, u, h, insecure)
|
||||
if err != nil {
|
||||
logger.Warn(fmt.Sprintf("%d error: %v", jobIndex, err))
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
}
|
||||
wg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateUrl(u string) string {
|
||||
// Parse the URL
|
||||
parsedURL, err := url.Parse(u)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error parsing URL: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Ensure it's a gemini URL
|
||||
if parsedURL.Scheme != "gemini" {
|
||||
fmt.Fprintf(os.Stderr, "Error: URL must use gemini:// scheme\n")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Get host and port
|
||||
host := parsedURL.Host
|
||||
if !strings.Contains(host, ":") {
|
||||
host = host + ":1965" // Default Gemini port
|
||||
}
|
||||
|
||||
return host
|
||||
}
|
||||
|
||||
func connect(ctx context.Context, url string, host string, insecure bool) error {
|
||||
logger := logging.FromContext(ctx)
|
||||
tlsConfig := &tls.Config{
|
||||
InsecureSkipVerify: insecure,
|
||||
MinVersion: tls.VersionTLS12,
|
||||
}
|
||||
|
||||
// Context checkpoint
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Connect to the server
|
||||
conn, err := tls.Dial("tcp", host, tlsConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Set connection deadline based on context
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
_ = conn.SetDeadline(deadline)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
// Context checkpoint
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Send the request (URL + CRLF)
|
||||
request := url + "\r\n"
|
||||
_, err = conn.Write([]byte(request))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Context checkpoint
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Read and dump response
|
||||
_, err = io.Copy(io.Discard, conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
jobIndex := ctx.Value(ctxKeyJobIndex)
|
||||
logger.Debug(fmt.Sprintf("%d done", jobIndex))
|
||||
return nil
|
||||
}
|
||||
@@ -22,7 +22,7 @@ import (
|
||||
|
||||
func main() {
|
||||
config.CONFIG = *config.GetConfig()
|
||||
logging.SetupLogging()
|
||||
logging.SetupLogging(config.CONFIG.LogLevel)
|
||||
logger := logging.Logger
|
||||
ctx := logging.WithLogger(context.Background(), logger)
|
||||
err := runApp(ctx)
|
||||
@@ -42,9 +42,11 @@ func runApp(ctx context.Context) error {
|
||||
signals := make(chan os.Signal, 1)
|
||||
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
|
||||
|
||||
// Only this file should send to this channel.
|
||||
// All external functions should return errors.
|
||||
fatalErrors := make(chan error)
|
||||
|
||||
// Root server context, used to cancel
|
||||
// Root context, used to cancel
|
||||
// connections and graceful shutdown.
|
||||
serverCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
Reference in New Issue
Block a user