Skip to content

demos/counter: introduce server and improve code readability #1505

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Jaana Dogan (rakyll) wants to merge 1 commit into from
Closed

demos/counter: introduce server and improve code readability #1505

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
331 changes: 181 additions & 150 deletions demos/counter/counter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,185 +38,218 @@ import (
"github.com/spf13/pflag"
)

var (
requestCount uint64
ready atomic.Bool
fileMutex sync.Mutex
sigtermSleepDurationSecs atomic.Int64
)
func main() {
counterDir := pflag.String("file-counter-directory", "/home/counter", "Directory for file counter")
secondCounterDir := pflag.String("second-file-counter-directory", "", "Directory for a second file counter; empty disables it. Used to exercise an Actor with more than one durable volume")
validateExistingFilePath := pflag.String("validate-existing-file-path", "", "Path to an existing file to validate reading; empty disables it")
extraPort := pflag.Int("extra-port", 0, "Additional port to listen on to test atenet-router arbitrary-port ingress; 0 disables it")
tcpPort := pflag.Int("tcp-port", 0, "Port for TCP echo to test atunnel CONNECT ingress; 0 disables it")
pflag.Parse()

ctx := context.Background()
slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil)))

srv := newServer(*counterDir, *secondCounterDir, *validateExistingFilePath)
srv.handleSignals(ctx)

go startHTTPServer(ctx, ":80", srv, "Starting counter server on port 80")

if *extraPort > 0 {
go startExtraPortServer(ctx, *extraPort)
}

if *tcpPort > 0 {
go startTCPEchoServer(ctx, *tcpPort)
}

// Write random data to a file in the root filesystem to test checkpoint/restore.
if err := writeRandomFile(); err != nil {
slog.InfoContext(ctx, "Error writing random file", slog.Any("err", err))
} else {
slog.InfoContext(ctx, "Wrote content to random file", slog.String("fshash", hashRandomFile()))
}

srv.setReady()
slog.InfoContext(ctx, "Readyz now reports OK")

logPeriodically(ctx)
}

func incrementFileCounter(filePath string) int {
fileMutex.Lock()
defer fileMutex.Unlock()
counter := 0
data, err := os.ReadFile(filePath)
if err == nil {
type server struct {
mux *http.ServeMux

counterDir string
secondCounterDir string
validateExistingFilePath string

requestCount atomic.Uint64
ready atomic.Bool

shutdownDelaySecs atomic.Int64

fileMu sync.Mutex // guards file operations
}

func newServer(counterDir, secondCounterDir, validateExistingFilePath string) *server {
s := &server{
mux: http.NewServeMux(),
counterDir: counterDir,
secondCounterDir: secondCounterDir,
validateExistingFilePath: validateExistingFilePath,
}
s.mux.HandleFunc("/", s.handle)
s.mux.HandleFunc("/readyz", s.handleReadyz)
s.mux.HandleFunc("/set-sigterm-sleep", s.storeShutdownDelay)
s.shutdownDelaySecs.Store(15)
return s
}

func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
s.mux.ServeHTTP(w, r)
}

func (s *server) setReady() {
s.ready.Store(true)
}

func (s *server) increment(path string) int {
s.fileMu.Lock()
defer s.fileMu.Unlock()

var counter int
if data, err := os.ReadFile(path); err == nil {
if i, err := strconv.Atoi(string(data)); err == nil {
counter = i
}
}
counter++
err = os.WriteFile(filePath, []byte(strconv.Itoa(counter)), 0o644)
if err != nil {

if err := os.WriteFile(path, []byte(strconv.Itoa(counter)), 0o644); err != nil {
return -1
}
return counter
}

func main() {
sigtermSleepDurationSecs.Store(15)
fileCounterDirectory := pflag.String("file-counter-directory", "/home/counter", "Directory for file counter")
secondFileCounterDirectory := pflag.String("second-file-counter-directory", "", "Directory for a second file counter; empty disables it. Used to exercise an Actor with more than one durable volume")
validateExistingFilePath := pflag.String("validate-existing-file-path", "", "Path to existing file to validate reading")
extraPort := pflag.Int("extra-port", 0, "Additional port to listen on, for exercising atenet-router's arbitrary-port ingress support; 0 disables it")
tcpPort := pflag.Int("tcp-port", 0, "Plain TCP echo port for exercising atunnel CONNECT ingress; 0 disables it")
pflag.Parse()
ctx := context.Background()

sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGTERM)
func (s *server) handleSignals(ctx context.Context) {
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGTERM)
go func() {
sig := <-sigCh
slog.InfoContext(ctx, "Received signal, waiting before exiting", slog.String("signal", sig.String()), slog.Int64("sleep_secs", sigtermSleepDurationSecs.Load()))
time.Sleep(time.Duration(sigtermSleepDurationSecs.Load()) * time.Second)
sig := <-ch
secs := s.shutdownDelaySecs.Load()
slog.InfoContext(ctx, "Received signal, waiting before exiting", slog.String("signal", sig.String()), slog.Int64("sleep_secs", secs))
time.Sleep(time.Duration(secs) * time.Second)
slog.InfoContext(ctx, "Exiting now")
os.Exit(0)
}()
}

slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil)))
func (s *server) handle(w http.ResponseWriter, r *http.Request) {
fileCounter := s.increment(filepath.Join(s.counterDir, "a.txt"))
memoryCounter := s.requestCount.Add(1)
currentIP := resolveCurrentIP()

defaultMux := http.NewServeMux()
defaultMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
fileCounter := incrementFileCounter(filepath.Join(*fileCounterDirectory, "a.txt"))
memoryCounter := atomic.AddUint64(&requestCount, 1)
currentIP := getCurrentIP()

fileContentStr := ""
if *validateExistingFilePath != "" {
fileContent, err := os.ReadFile(*validateExistingFilePath)
if err != nil {
fileResponse := fmt.Sprintf("failed to read test file: %s\n", err.Error())
w.WriteHeader(http.StatusOK)
w.Write([]byte(fileResponse))
return
}
fileContentStr = fmt.Sprintf(" | file content: %s", string(fileContent))
var content string
if s.validateExistingFilePath != "" {
fileContent, err := os.ReadFile(s.validateExistingFilePath)
if err != nil {
fmt.Fprintf(w, "failed to read test file: %s\n", err)
return
}
content = fmt.Sprintf(" | file content: %s", string(fileContent))
}

// A second counter in another directory, so an Actor with more than one
// durable volume can show each of them persisting independently.
secondFileCounterStr := ""
if *secondFileCounterDirectory != "" {
secondFileCounter := incrementFileCounter(filepath.Join(*secondFileCounterDirectory, "a.txt"))
secondFileCounterStr = fmt.Sprintf(" | preserved second file counter: %d", secondFileCounter)
}
// Optional second counter in another directory for multi-volume persistence test.
var secondContent string
if s.secondCounterDir != "" {
secondCounter := s.increment(filepath.Join(s.secondCounterDir, "a.txt"))
secondContent = fmt.Sprintf(" | preserved second file counter: %d", secondCounter)
}

response := fmt.Sprintf("hello from: %s | preserved memory count: %d | preserved file counter: %d%s%s\n", currentIP, memoryCounter, fileCounter, secondFileCounterStr, fileContentStr)
slog.InfoContext(ctx, "Handled request", slog.String("response", response))
body := fmt.Sprintf("hello from: %s | preserved memory count: %d | preserved file counter: %d%s%s\n",
currentIP, memoryCounter, fileCounter, secondContent, content)
slog.InfoContext(r.Context(), "Handled request", slog.String("body", body))

w.WriteHeader(http.StatusOK)
w.Write([]byte(response))
})
// /readyz is the endpoint the ateom-gvisor readyz probe polls. It returns
// 200 only once initialization (the random-file write) has completed.
// After a checkpoint+restore the atomic flag is part of the snapshot, so
// the endpoint returns 200 immediately on resume.
defaultMux.HandleFunc("/readyz", func(w http.ResponseWriter, _ *http.Request) {
if !ready.Load() {
http.Error(w, "not ready", http.StatusServiceUnavailable)
return
}
w.WriteHeader(http.StatusOK)
w.Write([]byte("ok\n"))
})
fmt.Fprint(w, body)
}

defaultMux.HandleFunc("/set-sigterm-sleep", func(w http.ResponseWriter, r *http.Request) {
durationStr := r.URL.Query().Get("duration")
if durationStr == "" {
http.Error(w, "missing duration parameter", http.StatusBadRequest)
return
}
d, err := strconv.Atoi(durationStr)
if err != nil || d < 0 {
http.Error(w, "invalid duration parameter", http.StatusBadRequest)
return
}
sigtermSleepDurationSecs.Store(int64(d))
response := fmt.Sprintf("SIGTERM sleep duration set to %d seconds\n", d)
slog.InfoContext(r.Context(), "Updated SIGTERM sleep duration", slog.Int("duration_secs", d))
w.WriteHeader(http.StatusOK)
w.Write([]byte(response))
// handleReadyz is polled by the ateom-gvisor ready probe.
func (s *server) handleReadyz(w http.ResponseWriter, _ *http.Request) {
if !s.ready.Load() {
http.Error(w, "not ready", http.StatusServiceUnavailable)
return
}
fmt.Fprint(w, "ok\n")
}

func (s *server) storeShutdownDelay(w http.ResponseWriter, r *http.Request) {
raw := r.URL.Query().Get("duration")
if raw == "" {
http.Error(w, "missing duration parameter", http.StatusBadRequest)
return
}
secs, err := strconv.ParseInt(raw, 10, 64)
if err != nil || secs < 0 {
http.Error(w, "invalid duration parameter", http.StatusBadRequest)
return
}
s.shutdownDelaySecs.Store(secs)
slog.InfoContext(r.Context(), "Updated SIGTERM sleep duration", slog.Int64("duration_secs", secs))
fmt.Fprintf(w, "SIGTERM sleep duration set to %d seconds\n", secs)
}

func startHTTPServer(ctx context.Context, addr string, handler http.Handler, msg string) {
slog.InfoContext(ctx, msg)
if err := http.ListenAndServe(addr, handler); err != nil {
slog.ErrorContext(ctx, "Error starting HTTP server", slog.String("addr", addr), slog.Any("err", err))
os.Exit(1)
}
}

func startExtraPortServer(ctx context.Context, port int) {
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
body := fmt.Sprintf("hello from extra port %d on pod %s\n", port, resolveCurrentIP())
slog.InfoContext(r.Context(), "Handled extra-port request", slog.String("body", body))
fmt.Fprint(w, body)
})

go func() {
slog.InfoContext(ctx, "Starting counter server on port 80")
if err := http.ListenAndServe(":80", defaultMux); err != nil {
slog.ErrorContext(ctx, "Error starting server", slog.Any("err", err))
os.Exit(1)
}
}()
addr := fmt.Sprintf(":%d", port)
startHTTPServer(ctx, addr, mux, fmt.Sprintf("Starting counter extra-port server on port %d", port))
}

// A second, independent listener a test can address to prove traffic
// actually reached this port rather than falling through to the default
// one -- see atenet-router's arbitrary-port ingress support.
if *extraPort > 0 {
extraMux := http.NewServeMux()
extraMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
response := fmt.Sprintf("hello from extra port %d on pod %s\n", *extraPort, getCurrentIP())
slog.InfoContext(r.Context(), "Handled extra-port request", slog.String("response", response))
w.WriteHeader(http.StatusOK)
w.Write([]byte(response))
})
go func() {
addr := fmt.Sprintf(":%d", *extraPort)
slog.InfoContext(ctx, "Starting counter extra-port server", slog.Int("port", *extraPort))
if err := http.ListenAndServe(addr, extraMux); err != nil {
slog.ErrorContext(ctx, "Error starting extra-port server", slog.Any("err", err))
os.Exit(1)
}
}()
func startTCPEchoServer(ctx context.Context, port int) {
addr := fmt.Sprintf(":%d", port)
listener, err := net.Listen("tcp", addr)
if err != nil {
slog.ErrorContext(ctx, "Error starting counter TCP echo server", slog.Int("port", port), slog.Any("err", err))
os.Exit(1)
}
defer listener.Close()
slog.InfoContext(ctx, "Starting counter TCP echo server", slog.Int("port", port))

if *tcpPort > 0 {
for {
conn, err := listener.Accept()
if err != nil {
slog.ErrorContext(ctx, "Counter TCP echo accept failed", slog.Any("err", err))
return
}
go func() {
listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *tcpPort))
if err != nil {
slog.ErrorContext(ctx, "Error starting counter TCP echo server", slog.Any("err", err))
os.Exit(1)
}
slog.InfoContext(ctx, "Starting counter TCP echo server", slog.Int("port", *tcpPort))
for {
conn, err := listener.Accept()
if err != nil {
slog.ErrorContext(ctx, "Counter TCP echo accept failed", slog.Any("err", err))
return
}
go func() {
defer conn.Close()
_, _ = io.Copy(conn, conn)
}()
}
defer conn.Close()
_, _ = io.Copy(conn, conn)
}()
}
}

// Write some random data to a file in the root filesystem, to test
// filesystem checkpoint/restore.
if err := writeRandomFile(); err != nil {
slog.InfoContext(ctx, "Error writing random file", slog.Any("err", err))
} else {
slog.InfoContext(ctx, "Wrote content to random file", slog.String("fshash", hashRandomFile()))
}

ready.Store(true)
slog.InfoContext(ctx, "Readyz now reports OK")

func logPeriodically(ctx context.Context) {
count := 0
slog.InfoContext(ctx, "Count", slog.Int("count", count), slog.String("fshash", hashRandomFile()))
count++

for range time.Tick(10 * time.Second) {
// TODO: Test outbound connectivity by pinging google.com
ticker := time.NewTicker(10 * time.Second)
defer ticker.Stop()

for range ticker.C {
// TODO(liorlieberman): Test outbound connectivity by pinging google.com
slog.InfoContext(ctx, "Count", slog.Int("count", count), slog.String("fshash", hashRandomFile()))
count++
}
Expand All @@ -225,15 +258,13 @@ func main() {
func writeRandomFile() error {
rf, err := os.Create("/random-content-file")
if err != nil {
return fmt.Errorf("while opening file: %w", err)
return fmt.Errorf("opening file: %w", err)
}
defer rf.Close()

_, err = io.CopyN(rf, rand.Reader, 1*1024*1024)
if err != nil {
return fmt.Errorf("while copying rand data: %w", err)
if _, err := io.CopyN(rf, rand.Reader, 1*1024*1024); err != nil {
return fmt.Errorf("copying random data: %w", err)
}

return nil
}

Expand All @@ -247,7 +278,7 @@ func hashRandomFile() string {
return base64.RawStdEncoding.EncodeToString(hash[:])
}

func getCurrentIP() string {
func resolveCurrentIP() string {
addrs, err := net.InterfaceAddrs()
if err != nil {
slog.Error("Error getting interface addresses", slog.Any("err", err))
Expand Down
Loading