143 lines
3.7 KiB
Go
143 lines
3.7 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/tls"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type settings struct {
|
|
Workers int `json:"workers"`
|
|
Connections int `json:"connections"`
|
|
TimeoutMS int `json:"timeout_ms"`
|
|
LifetimeMS int `json:"lifetime_ms"`
|
|
Certificate string `json:"cert_path"`
|
|
Key string `json:"key_path"`
|
|
}
|
|
|
|
type limitedListener struct {
|
|
net.Listener
|
|
slots chan struct{}
|
|
closed chan struct{}
|
|
once sync.Once
|
|
}
|
|
|
|
type limitedConnection struct {
|
|
net.Conn
|
|
slots chan struct{}
|
|
once sync.Once
|
|
}
|
|
|
|
func (connection *limitedConnection) Close() error {
|
|
err := connection.Conn.Close()
|
|
connection.once.Do(func() { <-connection.slots })
|
|
return err
|
|
}
|
|
|
|
func (listener *limitedListener) Accept() (net.Conn, error) {
|
|
connection, err := listener.Listener.Accept()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
select {
|
|
case listener.slots <- struct{}{}:
|
|
return &limitedConnection{Conn: connection, slots: listener.slots}, nil
|
|
case <-listener.closed:
|
|
connection.Close()
|
|
return nil, net.ErrClosed
|
|
}
|
|
}
|
|
|
|
func (listener *limitedListener) Close() error {
|
|
listener.once.Do(func() { close(listener.closed) })
|
|
return listener.Listener.Close()
|
|
}
|
|
|
|
func run() error {
|
|
if len(os.Args) != 2 || (os.Args[1] != "http" && os.Args[1] != "https") {
|
|
return errors.New("usage: go_http http|https")
|
|
}
|
|
executable, err := os.Executable()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
content, err := os.ReadFile(filepath.Join(filepath.Dir(executable), "config", "setting.json"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var document struct {
|
|
Benchmark settings `json:"benchmark"`
|
|
}
|
|
if err := json.Unmarshal(content, &document); err != nil {
|
|
return err
|
|
}
|
|
config := document.Benchmark
|
|
if config.Workers < 1 || config.Connections < 1 || config.TimeoutMS < 1 || config.LifetimeMS < 1 {
|
|
return errors.New("benchmark requires positive workers, connections, timeout_ms and lifetime_ms")
|
|
}
|
|
runtime.GOMAXPROCS(config.Workers)
|
|
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
bounded := &limitedListener{Listener: listener, slots: make(chan struct{}, config.Connections), closed: make(chan struct{})}
|
|
defer bounded.Close()
|
|
large := bytes.Repeat([]byte{'x'}, 65536)
|
|
deadline := time.Duration(config.TimeoutMS) * time.Millisecond
|
|
server := &http.Server{
|
|
ReadTimeout: deadline, ReadHeaderTimeout: deadline, WriteTimeout: deadline,
|
|
IdleTimeout: deadline, MaxHeaderBytes: 8192,
|
|
TLSNextProto: map[string]func(*http.Server, *tls.Conn, http.Handler){},
|
|
Handler: http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
|
body := http.MaxBytesReader(writer, request.Body, 65536)
|
|
defer body.Close()
|
|
if _, err := io.Copy(io.Discard, body); err != nil {
|
|
http.Error(writer, "invalid request body", http.StatusBadRequest)
|
|
return
|
|
}
|
|
switch request.URL.RequestURI() {
|
|
case "/plain":
|
|
writer.Header().Set("Content-Type", "text/plain")
|
|
io.WriteString(writer, "hello")
|
|
case "/json":
|
|
writer.Header().Set("Content-Type", "application/json")
|
|
io.WriteString(writer, `{"ok":true}`)
|
|
case "/bytes":
|
|
writer.Header().Set("Content-Type", "application/octet-stream")
|
|
writer.Write(large)
|
|
default:
|
|
http.NotFound(writer, request)
|
|
}
|
|
}),
|
|
}
|
|
timer := time.AfterFunc(time.Duration(config.LifetimeMS)*time.Millisecond, func() { server.Close() })
|
|
defer timer.Stop()
|
|
fmt.Printf("READY|%d\n", listener.Addr().(*net.TCPAddr).Port)
|
|
if os.Args[1] == "https" {
|
|
err = server.ServeTLS(bounded, config.Certificate, config.Key)
|
|
} else {
|
|
err = server.Serve(bounded)
|
|
}
|
|
if errors.Is(err, http.ErrServerClosed) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
|
|
func main() {
|
|
if err := run(); err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(1)
|
|
}
|
|
}
|