country-geo-locations/internal/exporter/middleware.go

68 lines
1.2 KiB
Go

package exporter
import (
"net/http"
"sync"
"time"
"github.com/go-chi/chi/v5/middleware"
)
type RequestData struct {
Latency time.Duration
Request *http.Request
Start time.Time
}
type RequestDataQueue struct {
mu sync.Mutex
data []RequestData
}
func NewRequestDataQueue() *RequestDataQueue {
return &RequestDataQueue{
data: []RequestData{},
}
}
func (q *RequestDataQueue) Add(data RequestData) {
q.mu.Lock()
defer q.mu.Unlock()
q.data = append(q.data, data)
}
func (q *RequestDataQueue) ConsumeAll() []RequestData {
q.mu.Lock()
defer q.mu.Unlock()
data := q.data
q.data = nil
return data
}
type Middleware struct {
queue *RequestDataQueue
}
func NewMiddleware(queue *RequestDataQueue) func(next http.Handler) http.Handler {
m := Middleware{
queue: queue,
}
return m.handler
}
func (m Middleware) handler(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
next.ServeHTTP(ww, r)
m.queue.Add(
RequestData{
Latency: time.Since(start),
Request: r,
Start: start,
},
)
}
return http.HandlerFunc(fn)
}