1/*
2 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
3 *
4 * Permission to use, copy, modify, and distribute this software for any
5 * purpose with or without fee is hereby granted, provided that the above
6 * copyright notice and this permission notice appear in all copies.
7 *
8 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15 */
16
17package miniprofiler_gae
18
19import (
20	"fmt"
21	"net/http"
22
23	"appengine"
24	"appengine/memcache"
25	"appengine/user"
26	"appengine_internal"
27	"github.com/MiniProfiler/go/miniprofiler"
28	"github.com/mjibson/appstats"
29)
30
31func init() {
32	miniprofiler.Enable = EnableIfAdminOrDev
33	miniprofiler.Get = GetMemcache
34	miniprofiler.Store = StoreMemcache
35	miniprofiler.MachineName = Instance
36}
37
38// EnableIfAdminOrDev returns true if this is the dev server or the current
39// user is an admin. This is the default for miniprofiler.Enable.
40func EnableIfAdminOrDev(r *http.Request) bool {
41	if appengine.IsDevAppServer() {
42		return true
43	}
44	c := appengine.NewContext(r)
45	if u := user.Current(c); u != nil {
46		return u.Admin
47	}
48	return false
49}
50
51// Instance returns the app engine instance id, or the hostname on dev.
52// This is the default for miniprofiler.MachineName.
53func Instance() string {
54	if i := appengine.InstanceID(); i != "" && !appengine.IsDevAppServer() {
55		return i[len(i)-8:]
56	}
57	return miniprofiler.Hostname()
58}
59
60// StoreMemcache stores the Profile in memcache. This is the default for
61// miniprofiler.Store.
62func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
63	item := &memcache.Item{
64		Key:   mp_key(string(p.Id)),
65		Value: p.Json(),
66	}
67	c := appengine.NewContext(r)
68	memcache.Set(c, item)
69}
70
71// GetMemcache gets the Profile from memcache. This is the default for
72// miniprofiler.Get.
73func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
74	c := appengine.NewContext(r)
75	item, err := memcache.Get(c, mp_key(id))
76	if err != nil {
77		return nil
78	}
79	return miniprofiler.ProfileFromJson(item.Value)
80}
81
82type Context struct {
83	appstats.Context
84	miniprofiler.Timer
85}
86
87func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) (err error) {
88	if c.Timer != nil && service != "__go__" {
89		c.StepCustomTiming(service, method, fmt.Sprintf("%v\n\n%v", method, in.String()), func() {
90			err = c.Context.Call(service, method, in, out, opts)
91		})
92	} else {
93		err = c.Context.Call(service, method, in, out, opts)
94	}
95	return
96}
97
98func (c Context) Step(name string, f func(Context)) {
99	if c.Timer != nil {
100		c.Timer.Step(name, func(t miniprofiler.Timer) {
101			f(Context{
102				Context: c.Context,
103				Timer:   t,
104			})
105		})
106	} else {
107		f(c)
108	}
109}
110
111// NewHandler returns a profiled, appstats-aware appengine.Context.
112func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) http.Handler {
113	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
114		h := miniprofiler.NewHandler(func(t miniprofiler.Timer, w http.ResponseWriter, r *http.Request) {
115			pc := Context{
116				Context: c.(appstats.Context),
117				Timer:   t,
118			}
119			t.SetName(miniprofiler.FuncName(f))
120			f(pc, w, r)
121			t.AddCustomLink("appstats", pc.URL())
122		})
123		h.ServeHTTP(w, r)
124	})
125}
126
127func mp_key(id string) string {
128	return fmt.Sprintf("mini-profiler-results:%s", id)
129}
130