grace timeouts from cfg

init.go

@@ -30,6 +30,12 @@ func init() {
 	flag.BoolVar(&cfgtestFlag, "t", false, "config test")
 }
 
+type Config struct {
+	Timeout      time.Duration
+	ReadTimeout  time.Duration
+	WriteTimeout time.Duration
+}
+
 // applications need some way to access the port
 // TODO: this method will work only after grace.Serve is called.
 func GetListenPort(hport string) string {
@@ -82,7 +88,7 @@ func ServerFastHTTP(hport string, handler fasthttp.RequestHandler) error {
 
 // start serving on hport. If running via socketmaster, the hport argument is
 // ignored. Also, if a port was specified via -p, it takes precedence on hport
-func Serve(hport string, handler http.Handler) error {
+func Serve(hport string, handler http.Handler, cfg *Config) error {
 
 	checkConfigTest()
 
@@ -92,11 +98,11 @@ func Serve(hport string, handler http.Handler) error {
 	}
 
 	srv := &graceful.Server{
-		Timeout: 10 * time.Second,
+		Timeout: cfg.Timeout * time.Second,
 		Server: &http.Server{
 			Handler:      handler,
-			ReadTimeout:  5 * time.Second,
-			WriteTimeout: 10 * time.Second,
+			ReadTimeout:  cfg.ReadTimeout * time.Second,
+			WriteTimeout: cfg.WriteTimeout * time.Second,
 		},
 	}
 

init_test.go

@@ -1,15 +1,23 @@
 package grace
 
 import (
-  "net/http"
-  "log"
+	"log"
+	"net/http"
+	"time"
 )
 
 func ExampleServe() {
-  http.HandleFunc("/foo/bar", foobarHandler)
-  log.Fatal(grace.Serve(":9000", nil))
+	http.HandleFunc("/foo/bar", foobarHandler)
+	cfg := &Config{
+		Timeout:      5 * time.Second,
+		ReadTimeout:  5 * time.Second,
+		WriteTimeout: 10 * time.Second,
+	}
+	// log.Fatal(grace.Serve(":9000", nil, cfg))
+	log.Fatal(Serve(":9000", nil, cfg))
+
 }
 
 func foobarHandler(w http.ResponseWriter, r *http.Request) {
-  w.Write([]byte("foobar"))
+	w.Write([]byte("foobar"))
 }