initial commit of api

2020-08-06 14:15:14 -04:00
parent 5db9d387df
commit a651b9cd73
9 changed files with 420 additions and 0 deletions
@@ -0,0 +1,69 @@
+package http
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"math/rand"
+	"net/http"
+	"time"
+)
+
+func (s *Server) setLimit(w http.ResponseWriter, r *http.Request) {
+	var input struct {
+		Limit int `json:"limit"`
+	}
+
+	bodyBytes, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		s.respondError(w, http.StatusBadRequest, err)
+		return
+	}
+
+	err = json.Unmarshal(bodyBytes, &input)
+	if err != nil {
+		s.respondError(w, http.StatusBadRequest, err)
+	}
+
+	var rl RateLimiter
+	rl, err = NewMaxConcurrencyRateLimiter(input.Limit)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("setting new rate limiter: %+v\n", rl)
+	s.RateLimiter = rl
+}
+
+func (s *Server) getData(w http.ResponseWriter, r *http.Request) {
+	data := []struct {
+		Item1 string
+		Item2 string
+		Item3 string
+	}{
+		{
+			Item1: "books",
+			Item2: "hotdogs",
+			Item3: "trinkets",
+		},
+	}
+
+	bytes, err := json.Marshal(data)
+	if err != nil {
+		s.respondError(w, http.StatusInternalServerError, err)
+		return
+	}
+
+	// Sleep for a random amount of time up to 5 seconds.
+	time.Sleep(time.Duration(rand.Intn(500)) * time.Millisecond)
+
+	w.Write(bytes)
+
+	t := r.Context().Value("rate-limit-token").(*Token)
+	s.RateLimiter.Release(t)
+	return
+}
+
+func (s *Server) respondError(w http.ResponseWriter, statusCode int, err error) {
+	http.Error(w, err.Error(), statusCode)
+	return
+}
@@ -0,0 +1,29 @@
+package http
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"time"
+)
+
+type Server struct {
+	RateLimiter RateLimiter
+	server      *http.Server
+}
+
+func (s *Server) Run() error {
+	s.server = &http.Server{
+		Addr:         ":8080",
+		ReadTimeout:  30 * time.Second,
+		WriteTimeout: 30 * time.Second,
+		Handler:      s.BuildRoutes(),
+	}
+
+	fmt.Println("server listening on :8080")
+	return s.server.ListenAndServe()
+}
+
+func (s *Server) Shutdown(ctx context.Context) error {
+	return s.server.Shutdown(ctx)
+}
@@ -0,0 +1,174 @@
+package http
+
+import (
+	"errors"
+	"log"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+type RateLimiter interface {
+	Acquire() (*Token, error)
+	Release(*Token)
+	LimitExceededPerc() float32
+	ReleaseAll()
+}
+
+type Token struct {
+	ID        string
+	CreatedAt time.Time
+}
+
+func NewMaxConcurrencyRateLimiter(limit int) (RateLimiter, error) {
+	if limit <= 0 {
+		return nil, errors.New("invalid rate limit")
+	}
+
+	m := newManager(limit)
+
+	// Await.
+	func() {
+		go func() {
+			for {
+				select {
+				case <-m.inChan:
+					m.tryGenerateToken()
+				case t := <-m.releaseChan:
+					m.releaseToken(t)
+				}
+			}
+		}()
+	}()
+	return m, nil
+}
+
+func newManager(limit int) *Manager {
+	return &Manager{
+		errorChan:    make(chan error),
+		outChan:      make(chan *Token),
+		inChan:       make(chan struct{}),
+		activeTokens: make(map[string]*Token),
+		releaseChan:  make(chan *Token),
+		needToken:    0,
+		limit:        limit,
+		makeToken: func() *Token {
+			return &Token{
+				ID:        uuid.New().String(),
+				CreatedAt: time.Now().UTC(),
+			}
+		},
+	}
+}
+
+type Manager struct {
+	errorChan    chan error
+	releaseChan  chan *Token
+	outChan      chan *Token
+	inChan       chan struct{}
+	needToken    int64
+	activeTokens map[string]*Token
+	tokensMutex  sync.RWMutex
+	limit        int
+	makeToken    func() *Token
+}
+
+func (m *Manager) Acquire() (*Token, error) {
+	go func() {
+		m.inChan <- struct{}{}
+	}()
+
+	select {
+	case t := <-m.outChan:
+		return t, nil
+	case err := <-m.errorChan:
+		return nil, err
+	}
+}
+
+func (m *Manager) Release(t *Token) {
+	go func() {
+		m.releaseChan <- t
+	}()
+}
+
+func (m *Manager) ReleaseAll() {
+	m.tokensMutex.Lock()
+	for _, t := range m.activeTokens {
+		m.Release(t)
+	}
+	m.activeTokens = make(map[string]*Token)
+	m.tokensMutex.Unlock()
+}
+
+func (m *Manager) LimitExceededPerc() float32 {
+	return m.limitExceededPerc()
+}
+
+func (m *Manager) incNeedToken() {
+	atomic.AddInt64(&m.needToken, 1)
+}
+
+func (m *Manager) decNeedToken() {
+	atomic.AddInt64(&m.needToken, -1)
+}
+
+func (m *Manager) awaitingToken() bool {
+	return atomic.LoadInt64(&m.needToken) > 0
+}
+
+func (m *Manager) limitExceededPerc() float32 {
+	m.tokensMutex.RLock()
+	defer m.tokensMutex.RUnlock()
+	return float32(atomic.LoadInt64(&m.needToken)) / float32(m.limit)
+}
+
+func (m *Manager) isLimitExceeded() bool {
+	m.tokensMutex.RLock()
+	defer m.tokensMutex.RUnlock()
+
+	if len(m.activeTokens) >= m.limit {
+		return true
+	}
+	return false
+}
+
+func (m *Manager) releaseToken(t *Token) {
+	if t == nil {
+		log.Println("cannot release a nil token")
+		return
+	}
+
+	m.tokensMutex.Lock()
+	if _, ok := m.activeTokens[t.ID]; !ok {
+		log.Printf("unable to release token %s - not in use", t)
+		return
+	}
+
+	delete(m.activeTokens, t.ID)
+	m.tokensMutex.Unlock()
+
+	if m.awaitingToken() {
+		m.decNeedToken()
+		go m.tryGenerateToken()
+	}
+}
+
+func (m *Manager) tryGenerateToken() {
+	if m.isLimitExceeded() {
+		m.incNeedToken()
+		return
+	}
+
+	t := m.makeToken()
+
+	m.tokensMutex.Lock()
+	m.activeTokens[t.ID] = t
+	m.tokensMutex.Unlock()
+
+	go func() {
+		m.outChan <- t
+	}()
+}
@@ -0,0 +1,37 @@
+package http
+
+import (
+	"context"
+	"errors"
+	"math/rand"
+	"net/http"
+)
+
+func (s *Server) rateLimiter(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		token, err := s.RateLimiter.Acquire()
+		if err != nil {
+			panic(err)
+		}
+		ctx := context.WithValue(r.Context(), "rate-limit-token", token)
+		next.ServeHTTP(w, r.WithContext(ctx))
+	})
+}
+
+func (s *Server) faulty(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		rlExceedPerc := s.RateLimiter.LimitExceededPerc()
+		// Faults will scale to the percentage of requests that are above the limit.
+		// When 2x volume is exceeded all requests should fail.
+		if rlExceedPerc > 0.0 {
+			check := rand.Float32()
+			if rlExceedPerc > float32(check) {
+				s.respondError(w, http.StatusInternalServerError, errors.New("application error"))
+				// Transaction failed so release the token.
+				s.RateLimiter.Release(r.Context().Value("rate-limit-token").(*Token))
+				return
+			}
+		}
+		next.ServeHTTP(w, r)
+	})
+}
@@ -0,0 +1,24 @@
+package http
+
+import (
+	"net/http"
+
+	"github.com/go-chi/chi"
+	"github.com/go-chi/chi/middleware"
+)
+
+func (s *Server) BuildRoutes() http.Handler {
+	r := chi.NewRouter()
+	r.Use(middleware.Logger)
+
+	r.Post("/limit", s.setLimit)
+
+	r.Route("/data", func(r chi.Router) {
+		r.Use(s.rateLimiter)
+		r.Use(s.faulty)
+
+		r.Get("/", s.getData)
+	})
+
+	return r
+}
@@ -0,0 +1,46 @@
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/RedVentures/agentx-loadtest-lunchandlearn/cmd/api/http"
+)
+
+func main() {
+	rl, err := http.NewMaxConcurrencyRateLimiter(10)
+	if err != nil {
+		log.Fatal(err)
+	}
+	s := &http.Server{
+		RateLimiter: rl,
+	}
+
+	// Allow app to listen for OS Interrupts and SIGTERMS.
+	serverErrors := make(chan error, 1)
+	osSignals := make(chan os.Signal, 1)
+	signal.Notify(osSignals, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		serverErrors <- s.Run()
+	}()
+
+	// Handling for server errors and OS signals.
+	select {
+	case err := <-serverErrors:
+		log.Fatalf("error starting server: %v", err.Error())
+	case <-osSignals:
+		log.Println("starting server shutdown...")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		err := s.Shutdown(ctx)
+		if err != nil {
+			log.Fatalf("error trying to shutdown http server: %v", err.Error())
+		}
+	}
+}