From 48415d88c00182e6e489a93e704e9df0832368db Mon Sep 17 00:00:00 2001
From: Farzad Ghanei <644113+farzadghanei@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:39:52 +0200
Subject: [PATCH] Refactor running checks away from main, add shutdown after
 requests

- Moved the check-running code into separate functions for modular
  design and enhanced code resuability in 'chkok.go' and 'run_modes.go'
- Add option to shutdown the http server after a number of requests,
  useful for testing
---
 cmd/chkok.go            |  86 +-------------------------
 cmd/chkok_test.go       |  54 ++++++++++++++++
 conf.go                 |   3 +-
 conf_test.go            |   5 ++
 examples/config.yaml    |   1 +
 examples/test-http.yaml |  23 +++++++
 run_modes.go            | 132 ++++++++++++++++++++++++++++++++++++++++
 7 files changed, 219 insertions(+), 85 deletions(-)
 create mode 100644 examples/test-http.yaml
 create mode 100644 run_modes.go

diff --git a/cmd/chkok.go b/cmd/chkok.go
index 0a0fe4e..392ad5e 100644
--- a/cmd/chkok.go
+++ b/cmd/chkok.go
@@ -5,9 +5,7 @@ import (
 	"fmt"
 	"io"
 	"log"
-	"net/http"
 	"os"
-	"time"
 
 	"github.com/farzadghanei/chkok"
 )
@@ -45,87 +43,7 @@ func run(confPath, mode string, output io.Writer, verbose bool) int {
 		return chkok.ExConfig
 	}
 	if mode == ModeHTTP {
-		return runHTTP(&checkGroups, conf, logger)
+		return chkok.RunModeHTTP(&checkGroups, conf, logger)
 	}
-	return runCli(&checkGroups, conf, output, logger)
-}
-
-// run app in CLI mode, return exit code
-func runCli(checkGroups *chkok.CheckSuites, conf *chkok.Conf, output io.Writer, logger *log.Logger) int {
-	runner := chkok.Runner{Log: logger, Timeout: conf.Runners["default"].Timeout}
-	logger.Printf("running checks ...")
-	checks := runner.RunChecks(*checkGroups)
-	incompeleteChecks := 0
-	failed := 0
-	for _, chk := range checks {
-		if chk.Status() != chkok.StatusDone {
-			incompeleteChecks++
-		}
-		if !chk.Result().IsOK {
-			failed++
-		}
-		fmt.Fprintf(output, "check %s status %d ok: %v\n", chk.Name(), chk.Status(), chk.Result().IsOK)
-	}
-	if incompeleteChecks > 0 {
-		fmt.Fprintf(output, "%v checks didn't get to completion", incompeleteChecks)
-		return chkok.ExTempFail
-	}
-	if failed > 0 {
-		return chkok.ExSoftware
-	}
-	return chkok.ExOK
-}
-
-// run app in http server mode, return exit code
-func runHTTP(checkGroups *chkok.CheckSuites, conf *chkok.Conf, logger *log.Logger) int {
-	runner := chkok.Runner{Log: logger, Timeout: conf.Runners["default"].Timeout}
-	logger.Printf("running checks ...")
-
-	httpHandler := func(w http.ResponseWriter, r *http.Request) {
-		logger.Printf("http request: %v", r)
-		checks := runner.RunChecks(*checkGroups)
-		incompeleteChecks := 0
-		failed := 0
-		passed := 0
-		for _, chk := range checks {
-			if chk.Status() != chkok.StatusDone {
-				incompeleteChecks++
-			} else {
-				if chk.Result().IsOK {
-					passed++
-				} else {
-					failed++
-				}
-			}
-			logger.Printf("check %s status %d ok: %v", chk.Name(), chk.Status(), chk.Result().IsOK)
-		}
-		logger.Printf("Run checks done. passed: %v - failed: %v - timedout: %v", passed, failed, incompeleteChecks)
-		if incompeleteChecks > 0 {
-			w.WriteHeader(http.StatusGatewayTimeout) // 504
-			fmt.Fprintf(w, "TIMEDOUT")
-		} else if failed > 0 {
-			w.WriteHeader(http.StatusInternalServerError) // 500
-			fmt.Fprintf(w, "FAILED")
-		} else {
-			w.WriteHeader(http.StatusOK)
-			fmt.Fprintf(w, "OK")
-		}
-	}
-
-	http.HandleFunc("/", httpHandler)
-	// TODO: allow to set server timeouts from configuration
-	server := &http.Server{
-		Addr:         ":8080",
-		Handler:      nil, // use http.DefaultServeMux
-		ReadTimeout:  2 * time.Second,
-		WriteTimeout: 5 * time.Second,
-		IdleTimeout:  2 * time.Second,
-	}
-	logger.Printf("starting http server ...")
-	err := server.ListenAndServe()
-	if err != nil {
-		logger.Printf("http server failed to start: %v", err)
-		return chkok.ExSoftware
-	}
-	return chkok.ExOK
+	return chkok.RunModeCLI(&checkGroups, conf, output, logger)
 }
diff --git a/cmd/chkok_test.go b/cmd/chkok_test.go
index 0c012f7..e83693c 100644
--- a/cmd/chkok_test.go
+++ b/cmd/chkok_test.go
@@ -3,9 +3,17 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"context"
+	"io"
+	"net/http"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
+)
+
+const (
+	MaxHTTPRetries = 5
 )
 
 func TestRunCli(t *testing.T) {
@@ -19,3 +27,49 @@ func TestRunCli(t *testing.T) {
 		t.Errorf("want exit code 0, got %v. output: %v", got, buf.String())
 	}
 }
+
+func TestRunHttp(t *testing.T) {
+	var buf bytes.Buffer
+	cwd, _ := os.Getwd()
+	baseDir, _ := filepath.Abs(filepath.Dir(cwd))
+	writer := bufio.NewWriter(&buf)
+	var confPath = filepath.Join(baseDir, "examples", "test-http.yaml")
+
+	go func() { // run the server in a goroutine
+		run(confPath, ModeHTTP, writer, false)
+	}()
+
+	// Test the runner via an HTTP request
+	// Create a new request with a context
+	req, err := http.NewRequestWithContext(
+		context.Background(), http.MethodGet, "http://localhost:8080", http.NoBody)
+	if err != nil {
+		t.Fatalf("Failed to create HTTP request: %v", err)
+	}
+
+	// Send the request multiple times, waiting for the server to
+	// start up and respond
+	var resp *http.Response
+	var body []byte
+	for i := 0; i < MaxHTTPRetries; i++ {
+		resp, err = http.DefaultClient.Do(req)
+		if err == nil && resp != nil { // server is up
+			body, err = io.ReadAll(io.Reader(resp.Body))
+			resp.Body.Close()
+			if err != nil {
+				t.Fatalf("Failed to read HTTP runner response body: %v", err)
+			}
+			break
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+	if err != nil {
+		t.Fatalf("Failed to send HTTP request to HTTP runner after %v treis: %v", MaxHTTPRetries, err)
+	}
+
+	// Assert the response body
+	want := "OK"
+	if string(body) != want {
+		t.Errorf("want response body %q, got %q", want, body)
+	}
+}
diff --git a/conf.go b/conf.go
index 801f54a..a09cb69 100644
--- a/conf.go
+++ b/conf.go
@@ -21,7 +21,8 @@ type Conf struct {
 
 // ConfRunner is config for the check runners
 type ConfRunner struct {
-	Timeout time.Duration
+	Timeout               time.Duration
+	ShutdownAfterRequests uint32 `yaml:"shutdown_after_requests"`
 }
 
 // ConfCheckSpec is the spec for each check configuration
diff --git a/conf_test.go b/conf_test.go
index bd3e5eb..67b735b 100644
--- a/conf_test.go
+++ b/conf_test.go
@@ -36,6 +36,11 @@ func TestReadConf(t *testing.T) {
 	if runner.Timeout.Minutes() != float64(wantMinutes) {
 		t.Errorf("invalid read conf default runner, want %v timeout got %v", wantMinutes, runner.Timeout.Minutes())
 	}
+	var wantShutdownAfterRequests uint32 = 100
+	if runner.ShutdownAfterRequests != wantShutdownAfterRequests {
+		t.Errorf("invalid read conf, want %v shutdown after requests got %v",
+			wantShutdownAfterRequests, runner.ShutdownAfterRequests)
+	}
 	etcChecks, ok := conf.CheckSuites["etc"]
 	if !ok {
 		t.Errorf("read conf found no etc checks")
diff --git a/examples/config.yaml b/examples/config.yaml
index bd3f526..4e04b76 100644
--- a/examples/config.yaml
+++ b/examples/config.yaml
@@ -3,6 +3,7 @@
 runners:
   default:
     timeout: 5m
+    shutdown_after_requests: 100  # mainly useful for testing http mode
 
 check_suites:
   etc:
diff --git a/examples/test-http.yaml b/examples/test-http.yaml
new file mode 100644
index 0000000..663a522
--- /dev/null
+++ b/examples/test-http.yaml
@@ -0,0 +1,23 @@
+---
+# sample check definitions used as fixture in tests
+
+runners:
+  default:
+    timeout: 1m
+  http:
+    timeout: 5s
+    shutdown_after_requests: 1  # stop http server after 1 request, used for testing http mode
+
+check_suites:
+  default:
+    - type: dir
+      path: ../examples
+    - type: file
+      path: ../examples/test.yaml
+      min_size: 10
+  missing:
+    - type: file
+      path: doesnt/exist/but/ok
+      absent: true
+
+...
diff --git a/run_modes.go b/run_modes.go
new file mode 100644
index 0000000..c574f1a
--- /dev/null
+++ b/run_modes.go
@@ -0,0 +1,132 @@
+package chkok
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"sync/atomic"
+	"time"
+)
+
+const (
+	// ShutdownTimeout is the time to wait for server to shutdown
+	ShutdownTimeout = 5 * time.Second
+)
+
+// RunModeCLI run app in CLI mode using the provided configs, return exit code
+func RunModeCLI(checkGroups *CheckSuites, conf *Conf, output io.Writer, logger *log.Logger) int {
+	runner := Runner{Log: logger, Timeout: conf.Runners["default"].Timeout}
+	passed, failed, timedout := runChecks(&runner, checkGroups, logger)
+	total := passed + failed + timedout
+	if timedout > 0 {
+		fmt.Fprintf(output, "%v/%v checks timedout", timedout, total)
+		return ExTempFail
+	}
+	if failed > 0 {
+		fmt.Fprintf(output, "%v/%v checks failed", failed, total)
+		return ExSoftware
+	}
+	fmt.Fprintf(output, "%v checks passed", total)
+	return ExOK
+}
+
+func httpRequestAsString(r *http.Request) string {
+	return fmt.Sprintf("%s - %s %s %s", r.RemoteAddr, r.Proto, r.Method, r.URL)
+}
+
+// RunModeHTTP runs app in http server mode using the provided config, return exit code
+func RunModeHTTP(checkGroups *CheckSuites, conf *Conf, logger *log.Logger) int {
+	timeout := conf.Runners["default"].Timeout
+	shutdownAfterRequests := conf.Runners["default"].ShutdownAfterRequests
+
+	// override default runner config if with http runner config if provided
+	if httpRunnerConf, ok := conf.Runners["http"]; ok {
+		if httpRunnerConf.Timeout > 0 {
+			timeout = httpRunnerConf.Timeout
+		}
+		if httpRunnerConf.ShutdownAfterRequests > 0 {
+			shutdownAfterRequests = httpRunnerConf.ShutdownAfterRequests
+		}
+	}
+	runner := Runner{Log: logger, Timeout: timeout}
+
+	var reqHandlerChan = make(chan *http.Request, 1)
+
+	httpHandler := func(w http.ResponseWriter, r *http.Request) {
+		// TODO: custmize return codes and messages from configuration
+		logger.Printf("processing http request: %s", httpRequestAsString(r))
+		_, failed, timedout := runChecks(&runner, checkGroups, logger)
+		if timedout > 0 {
+			w.WriteHeader(http.StatusGatewayTimeout) // 504
+			fmt.Fprintf(w, "TIMEDOUT")
+		} else if failed > 0 {
+			w.WriteHeader(http.StatusInternalServerError) // 500
+			fmt.Fprintf(w, "FAILED")
+		} else {
+			w.WriteHeader(http.StatusOK)
+			fmt.Fprintf(w, "OK")
+		}
+		reqHandlerChan <- r
+	}
+
+	http.HandleFunc("/", httpHandler)
+	// TODO: allow to set server timeouts from configuration
+	server := &http.Server{
+		Addr:         ":8080",
+		Handler:      nil, // use http.DefaultServeMux
+		ReadTimeout:  2 * time.Second,
+		WriteTimeout: 5 * time.Second,
+		IdleTimeout:  2 * time.Second,
+	}
+
+	var count uint32 = 0
+
+	go func() {
+		var request *http.Request
+		timeoutCtx, cancel := context.WithTimeout(context.Background(), ShutdownTimeout)
+		defer cancel()
+		for request = range reqHandlerChan {
+			atomic.AddUint32(&count, 1)
+			logger.Printf("request [%v] is processed: %v", count, httpRequestAsString(request))
+			if shutdownAfterRequests > 0 && atomic.LoadUint32(&count) >= shutdownAfterRequests {
+				if err := server.Shutdown(timeoutCtx); err != nil {
+					logger.Printf("http server shutdown failed: %v", err)
+				}
+				return
+			}
+		}
+	}()
+
+	logger.Printf("starting http server ...")
+	err := server.ListenAndServe()
+	close(reqHandlerChan)
+	if err != nil {
+		if atomic.LoadUint32(&count) < 1 { // server didn't handle any requests
+			logger.Printf("http server failed to start: %v", err)
+			return ExSoftware
+		}
+	}
+	logger.Printf("http server shutdown!")
+	return ExOK
+}
+
+// runChecks runs checks with logs, and returns number of passed, failed and timedout checks
+func runChecks(runner *Runner, checkGroups *CheckSuites, logger *log.Logger) (passed, failed, timedout int) {
+	checks := runner.RunChecks(*checkGroups)
+	for _, chk := range checks {
+		if chk.Status() != StatusDone {
+			timedout++
+		} else {
+			if chk.Result().IsOK {
+				passed++
+			} else {
+				failed++
+			}
+		}
+		logger.Printf("check %s status %d ok: %v", chk.Name(), chk.Status(), chk.Result().IsOK)
+	}
+	logger.Printf("%v checks done. passed: %v - failed: %v - timedout: %v", len(checks), passed, failed, timedout)
+	return passed, failed, timedout
+}