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) } } }