1
0
mirror of https://github.com/Oxalide/vsphere-influxdb-go.git synced 2023-10-10 13:36:51 +02:00
vsphere-influxdb-go/vendor/github.com/influxdata/influxdb/services/httpd/requests.go
2017-10-25 20:52:40 +00:00

141 lines
2.7 KiB
Go

package httpd
import (
"container/list"
"fmt"
"net"
"net/http"
"sync"
"sync/atomic"
"github.com/influxdata/influxdb/services/meta"
)
type RequestInfo struct {
IPAddr string
Username string
}
type RequestStats struct {
Writes int64 `json:"writes"`
Queries int64 `json:"queries"`
}
func (r *RequestInfo) String() string {
if r.Username != "" {
return fmt.Sprintf("%s:%s", r.Username, r.IPAddr)
}
return r.IPAddr
}
type RequestProfile struct {
tracker *RequestTracker
elem *list.Element
mu sync.RWMutex
Requests map[RequestInfo]*RequestStats
}
func (p *RequestProfile) AddWrite(info RequestInfo) {
p.add(info, p.addWrite)
}
func (p *RequestProfile) AddQuery(info RequestInfo) {
p.add(info, p.addQuery)
}
func (p *RequestProfile) add(info RequestInfo, fn func(*RequestStats)) {
// Look for a request entry for this request.
p.mu.RLock()
st, ok := p.Requests[info]
p.mu.RUnlock()
if ok {
fn(st)
return
}
// There is no entry in the request tracker. Create one.
p.mu.Lock()
if st, ok := p.Requests[info]; ok {
// Something else created this entry while we were waiting for the lock.
p.mu.Unlock()
fn(st)
return
}
st = &RequestStats{}
p.Requests[info] = st
p.mu.Unlock()
fn(st)
}
func (p *RequestProfile) addWrite(st *RequestStats) {
atomic.AddInt64(&st.Writes, 1)
}
func (p *RequestProfile) addQuery(st *RequestStats) {
atomic.AddInt64(&st.Queries, 1)
}
// Stop informs the RequestTracker to stop collecting statistics for this
// profile.
func (p *RequestProfile) Stop() {
p.tracker.mu.Lock()
p.tracker.profiles.Remove(p.elem)
p.tracker.mu.Unlock()
}
type RequestTracker struct {
mu sync.RWMutex
profiles *list.List
}
func NewRequestTracker() *RequestTracker {
return &RequestTracker{
profiles: list.New(),
}
}
func (rt *RequestTracker) TrackRequests() *RequestProfile {
// Perform the memory allocation outside of the lock.
profile := &RequestProfile{
Requests: make(map[RequestInfo]*RequestStats),
tracker: rt,
}
rt.mu.Lock()
profile.elem = rt.profiles.PushBack(profile)
rt.mu.Unlock()
return profile
}
func (rt *RequestTracker) Add(req *http.Request, user meta.User) {
rt.mu.RLock()
if rt.profiles.Len() == 0 {
rt.mu.RUnlock()
return
}
defer rt.mu.RUnlock()
var info RequestInfo
host, _, err := net.SplitHostPort(req.RemoteAddr)
if err != nil {
return
}
info.IPAddr = host
if user != nil {
info.Username = user.ID()
}
// Add the request info to the profiles.
for p := rt.profiles.Front(); p != nil; p = p.Next() {
profile := p.Value.(*RequestProfile)
if req.URL.Path == "/query" {
profile.AddQuery(info)
} else if req.URL.Path == "/write" {
profile.AddWrite(info)
}
}
}