RateLimiter

stanley.go

package main

import (
	"fmt"
	"net/http"
	"sync"
)

// Max request limit.  Gives all customers the same limit.  Set to 1 so it is easy to test
// When the request count exceeds the limit a 429 is returned.
const maxlimit = 1
const limitExceeded = 429

type rateLimiter struct {
	mux    *http.ServeMux
	limits map[string]int
	lock   sync.Mutex
}

// Needs to be made robust.  Just a stub now
func uniqueid(r *http.Request) string {
	return "xyz"
}
func (m rateLimiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.lock.Lock()
	if limit, ok := m.limits[uniqueid(r)]; ok {
		if limit < maxlimit {

			m.limits[uniqueid(r)] = limit + 1
			m.lock.Unlock()

			m.mux.ServeHTTP(w, r)
			m.decrease(r)
		} else {
			m.lock.Unlock()
			w.WriteHeader(limitExceeded)

		}
	} else {
		if maxlimit > 0 {
			m.limits[uniqueid(r)] = 1
			m.lock.Unlock()
			m.mux.ServeHTTP(w, r)
			m.decrease(r)
		} else {
			m.lock.Unlock()
			w.WriteHeader(limitExceeded)
		}
	}
}

func (m rateLimiter) decrease(r *http.Request) {
	m.lock.Lock()
	limit := m.limits[uniqueid(r)]
	m.limits[uniqueid(r)] = limit - 1
	m.lock.Unlock()
}

func (m rateLimiter) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	m.mux.HandleFunc(pattern, handler)
}
func helloHandler(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintf(w, "hello\n")
}
func main() {
	rl := rateLimiter{
		mux:    http.NewServeMux(),
		limits: map[string]int{},
		lock:   sync.Mutex{}}
	rl.HandleFunc("/", helloHandler)
	http.ListenAndServe(":8080", rl)
}