Context with Timeout

package main

import (
	"context"
	"log"
	"net/http"
	"time"
	"fmt"
)

func fakeDbCall(confirmed chan int) {
	time.Sleep(1000 * time.Millisecond)
	confirmed <- 2
}

func AddTimeoutContext(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		//Add timeout to context
		ctx, _ := context.WithTimeout(r.Context(), 500*time.Millisecond)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func AddContext2(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Println(r.Method, "-", r.RequestURI)
		cookie, _ := r.Cookie("username")
		if cookie != nil {
			//Add data to context
			ctx := context.WithValue(r.Context(), "Username", cookie.Value)
			next.ServeHTTP(w, r.WithContext(ctx))
		} else {
			next.ServeHTTP(w, r)
		}
	})
}

func StatusPage2(w http.ResponseWriter, r *http.Request) {
	//Get data from context
	if username := r.Context().Value("Username"); username != nil {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Hello " + username.(string) + "\n"))
	} else {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("Not Logged in"))
	}
}

func DbPage(w http.ResponseWriter, r *http.Request) {
	confirmed := make(chan int)
	for {
		select {
		case <-confirmed:
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("Your Db operation has been completed successfully.\n"))
			return
		case <-r.Context().Done():
			if r.Context().Err() == context.Canceled {
				w.WriteHeader(http.StatusGatewayTimeout)
				w.Write([]byte("Your Db operation is canceled.\n"))
				fmt.Printf("Your Db operation is canceled.")
				return
			} else if r.Context().Err() == context.DeadlineExceeded {
				w.WriteHeader(http.StatusGatewayTimeout)
				w.Write([]byte("Your Db Operation timedout.\n"))
				return
			}
		default:
			go fakeDbCall(confirmed)
			time.Sleep(100*time.Millisecond)
		}
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Long operation completed fine.\n"))
}

func LoginPage2(w http.ResponseWriter, r *http.Request) {
	expiration := time.Now().Add(365 * 24 * time.Hour)  ////Set to expire in 1 year
	cookie := http.Cookie{Name: "username", Value: "[email protected]", Expires: expiration}
	http.SetCookie(w, &cookie)
}


func LogoutPage2(w http.ResponseWriter, r *http.Request) {
	expiration := time.Now().AddDate(0, 0, -1)	//Set to expire in the past
	cookie := http.Cookie{Name: "username", Value: "[email protected]", Expires: expiration}
	http.SetCookie(w, &cookie)
}


func main() {
	mux := http.NewServeMux()

	mux.HandleFunc("/", StatusPage2)
	mux.HandleFunc("/login", LoginPage2)
	mux.HandleFunc("/logout", LogoutPage2)
	mux.Handle("/db", AddTimeoutContext(http.HandlerFunc(DbPage)))

	log.Println("Start server on port :8085")
	contextedMux := AddContext2(mux)
	log.Fatal(http.ListenAndServe(":8085", contextedMux))
}