func (g *GrpcEndpoint) Allow(ctx context.Context, req *pb.AllowRequest) (*pb.AllowResponse, error) {
rsp := new(pb.AllowResponse)
if invalid(req) {
logging.Printf("Invalid request %+v", req)
rsp.Status = pb.AllowResponse_REJECTED_INVALID_REQUEST
return rsp, nil
}
var tokensRequested int64 = 1
if req.TokensRequested > 0 {
tokensRequested = req.TokensRequested
}
wait, err := g.qs.Allow(req.Namespace, req.BucketName, tokensRequested, req.MaxWaitMillisOverride, req.MaxWaitTimeOverride)
if err != nil {
if qsErr, ok := err.(quotaservice.QuotaServiceError); ok {
rsp.Status = toPBStatus(qsErr)
} else {
logging.Printf("Caught error %v", err)
rsp.Status = pb.AllowResponse_REJECTED_SERVER_ERROR
}
} else {
rsp.Status = pb.AllowResponse_OK
rsp.TokensGranted = req.TokensRequested
rsp.WaitMillis = wait.Nanoseconds() / int64(time.Millisecond)
}
return rsp, nil
}
func (h *uiHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if h.development {
if err := h.loadTemplates(); err != nil {
logging.Printf("Caught error %v reloading templates", err)
}
}
tpl := strings.TrimPrefix(r.URL.Path, "/admin/")
if tpl == "" {
tpl = "index.html"
}
if h.templates.Lookup(tpl) == nil {
http.NotFound(w, r)
return
}
err := h.templates.ExecuteTemplate(w, tpl, h.a.Configs())
if err != nil {
logging.Printf("Caught error %v serving URL %v", err, r.URL.Path)
http.Error(w, "500 internal server error", http.StatusInternalServerError)
}
}
func (bf *bucketFactory) connectToRedisLocked() {
// Set up connection to Redis
bf.client = redis.NewClient(bf.redisOpts)
redisResults := bf.client.Time().Val()
if len(redisResults) == 0 {
logging.Printf("Cannot connect to Redis. TIME returned %v", redisResults)
} else {
t := time.Unix(toInt64(redisResults[0], 0), 0)
logging.Printf("Connection established. Time on Redis server: %v", t)
}
bf.scriptSHA = loadScript(bf.client)
}
// HandleEvent is implemented for stats.Listener
// HandleEvent consumes dynamic bucket events (see events.Event)
func (l *redisListener) HandleEvent(event events.Event) {
if !event.Dynamic() {
return
}
var key string
var numTokens int64 = 1
switch event.EventType() {
case events.EVENT_BUCKET_MISS:
key = "misses"
case events.EVENT_TOKENS_SERVED:
numTokens = event.NumTokens()
key = "hits"
default:
return
}
namespace := statsNamespace(key, event.Namespace())
bucket := event.BucketName()
var incr *redis.FloatCmd
_, err := l.client.Pipelined(func(pipe *redis.Pipeline) error {
incr = pipe.ZIncrBy(namespace, float64(numTokens), bucket)
pipe.ExpireAt(namespace, nearestHour())
return nil
})
if err != nil || incr.Err() != nil {
logging.Printf("RedisStatsListener.HandleEvent error (%s, %s, %d) %v, %v",
namespace, bucket, numTokens, err, incr.Err())
}
}
func (r *ResponseWrapper) log() {
timeFormatted := r.time.Format("02/Jan/2006 03:04:05")
requestLine := fmt.Sprintf("%s %s %s", r.method, r.uri, r.protocol)
logging.Printf(LogPattern,
r.ip, timeFormatted, requestLine, r.status, r.responseBytes,
r.elapsedTime.Seconds())
}
func (g *GrpcEndpoint) Start() {
lis, err := net.Listen("tcp", g.hostport)
if err != nil {
logging.Fatalf("Cannot start server on port %v. Error %v", g.hostport, err)
panic(fmt.Sprintf("Cannot start server on port %v. Error %v", g.hostport, err))
}
grpclog.SetLogger(logging.CurrentLogger())
g.grpcServer = grpc.NewServer()
// Each service should be registered
pb.RegisterQuotaServiceServer(g.grpcServer, g)
go g.grpcServer.Serve(lis)
g.currentStatus = lifecycle.Started
logging.Printf("Starting server on %v", g.hostport)
logging.Printf("Server status: %v", g.currentStatus)
}
func (b *redisBucket) Take(requested int64, maxWaitTime time.Duration) (time.Duration, bool) {
currentTimeNanos := strconv.FormatInt(time.Now().UnixNano(), 10)
args := []string{currentTimeNanos, b.nanosBetweenTokens, b.maxTokensToAccumulate,
strconv.FormatInt(requested, 10), strconv.FormatInt(maxWaitTime.Nanoseconds(), 10),
b.maxIdleTimeMillis, b.maxDebtNanos}
keepTrying := true
var waitTime time.Duration
for attempt := 0; keepTrying && attempt < b.factory.connectionRetries; attempt++ {
client := b.factory.client
res := client.EvalSha(b.factory.scriptSHA, b.redisKeys, args)
switch waitTimeNanos := res.Val().(type) {
case int64:
waitTime = time.Nanosecond * time.Duration(waitTimeNanos)
keepTrying = false
default:
// Always close connections on errors to prevent results leaking.
if err := b.factory.client.Close(); err != nil {
logging.Printf("Received error on redis client close: %+v", err)
}
if res.Err() != nil && res.Err().Error() == "redis: client is closed" {
b.factory.reconnectToRedis(client)
} else {
logging.Printf("Unknown response '%v' of type %T. Full result %+v",
waitTimeNanos, waitTimeNanos, res)
b.factory.reconnectToRedis(client)
}
}
}
if keepTrying {
panic(fmt.Sprintf("Couldn't reconnect to Redis, even after %v attempts",
b.factory.connectionRetries))
}
if waitTime < 0 {
// Timed out
return 0, false
}
return waitTime, true
}
func writeJSONError(w http.ResponseWriter, err *HttpError) {
response := make(map[string]string)
response["error"] = http.StatusText(err.status)
response["description"] = err.message
logging.Printf("Response error: %+v", response)
w.WriteHeader(err.status)
writeJSON(w, response)
}
// loadScript loads the LUA script into Redis. The LUA script contains the token bucket algorithm
// which is executed atomically in Redis. Once the script is loaded, it is invoked using its SHA.
func loadScript(c *redis.Client) (sha string) {
lua := `
local tokensNextAvailableNanos = tonumber(redis.call("GET", KEYS[1]))
if not tokensNextAvailableNanos then
tokensNextAvailableNanos = 0
end
local maxTokensToAccumulate = tonumber(ARGV[3])
local accumulatedTokens = redis.call("GET", KEYS[2])
if not accumulatedTokens then
accumulatedTokens = maxTokensToAccumulate
end
local currentTimeNanos = tonumber(ARGV[1])
local nanosBetweenTokens = tonumber(ARGV[2])
local requested = tonumber(ARGV[4])
local maxWaitTime = tonumber(ARGV[5])
local lifespan = tonumber(ARGV[6])
local maxDebtNanos = tonumber(ARGV[7])
local freshTokens = 0
if currentTimeNanos > tokensNextAvailableNanos then
freshTokens = math.floor((currentTimeNanos - tokensNextAvailableNanos) / nanosBetweenTokens)
accumulatedTokens = math.min(maxTokensToAccumulate, accumulatedTokens + freshTokens)
tokensNextAvailableNanos = currentTimeNanos
end
local waitTime = tokensNextAvailableNanos - currentTimeNanos
local accumulatedTokensUsed = math.min(accumulatedTokens, requested)
local tokensToWaitFor = requested - accumulatedTokensUsed
local futureWaitNanos = tokensToWaitFor * nanosBetweenTokens
tokensNextAvailableNanos = tokensNextAvailableNanos + futureWaitNanos
accumulatedTokens = accumulatedTokens - accumulatedTokensUsed
if (tokensNextAvailableNanos - currentTimeNanos > maxDebtNanos) or (waitTime > 0 and waitTime > maxWaitTime) then
waitTime = -1
else
if lifespan > 0 then
redis.call("SET", KEYS[1], tokensNextAvailableNanos, "PX", lifespan)
redis.call("SET", KEYS[2], math.floor(accumulatedTokens), "PX", lifespan)
else
redis.call("SET", KEYS[1], tokensNextAvailableNanos)
redis.call("SET", KEYS[2], math.floor(accumulatedTokens))
end
end
return waitTime
`
s := c.ScriptLoad(lua)
sha = s.Val()
logging.Printf("Loaded LUA script into Redis; script SHA %v", sha)
return
}
func (z *ZkConfigPersister) currentConfigEventListener() (<-chan zk.Event, error) {
config, _, ch, err := z.conn.GetW(z.path)
if err != nil {
logging.Printf("Received error from zookeeper when fetching %s: %+v", z.path, err)
return nil, err
}
children, _, err := z.conn.Children(z.path)
if err != nil {
logging.Printf("Received error from zookeeper when fetching children of %s: %+v", z.path, err)
return nil, err
}
z.Lock()
defer z.Unlock()
z.configs = make(map[string][]byte)
for _, child := range children {
path := fmt.Sprintf("%s/%s", z.path, child)
data, _, err := z.conn.Get(path)
if err != nil {
logging.Printf("Received error from zookeeper when fetching %s: %+v", path, err)
} else {
z.configs[child] = data
}
}
z.config = string(config)
select {
case z.watcher <- struct{}{}:
// Notified
default:
// Doesn't matter; another notification is pending.
}
return ch, nil
}
func toInt64(s interface{}, defaultValue int64) int64 {
if s != nil {
v, err := strconv.ParseInt(s.(string), 10, 64)
if err != nil {
logging.Printf("Cannot convert '%v' to int64", s)
return defaultValue
}
return v
}
return defaultValue
}
// waitTimeLoop is the single event loop that claims tokens on a given bucket.
func (b *tokenBucket) waitTimeLoop() {
for {
select {
case req := <-b.waitTimer:
req.response <- b.calcWaitTime(req.requested, req.maxWaitTimeNanos)
case <-b.closer:
logging.Printf("Garbage collecting bucket %v", b.fullName)
// TODO(manik) properly notify goroutines who are currently trying to write to waitTimer
return
}
}
}
// Get is implemented for stats.Listener
// Get returns the hits and misses for a bucket in the specified namespace
// within the current bucketed hour
func (l *redisListener) Get(namespace, bucket string) *BucketScores {
scores := &BucketScores{0, 0}
value, err := l.client.ZScore(statsNamespace("misses", namespace), bucket).Result()
if err != nil && err.Error() != "redis: nil" {
logging.Printf("RedisStatsListener.Get error (%s, %s) %v", namespace, bucket, err)
} else {
scores.Misses = int64(value)
}
value, err = l.client.ZScore(statsNamespace("hits", namespace), bucket).Result()
if err != nil && err.Error() != "redis: nil" {
logging.Printf("RedisStatsListener.Get error (%s, %s) %v", namespace, bucket, err)
} else {
scores.Hits = int64(value)
}
return scores
}
func writeJSON(w http.ResponseWriter, object interface{}) {
b, e := json.Marshal(object)
if e != nil {
writeJSONError(w, &HttpError{e.Error(), http.StatusBadRequest})
return
}
_, e = w.Write(b)
if e != nil {
logging.Printf("Error writing JSON! %+v", e)
}
}
func (l *redisListener) redisTopList(key string) []*BucketScore {
results, err := l.client.ZRevRangeWithScores(key, 0, 10).Result()
if err != nil && err.Error() != "redis: nil" {
logging.Printf("RedisStatsListener.TopList error (%s) %v", key, err)
return emptyArr
}
arr := make([]*BucketScore, len(results))
for i, item := range results {
arr[i] = &BucketScore{item.Member.(string), int64(item.Score)}
}
return arr
}
func main() {
cfg := config.NewDefaultServiceConfig()
ns := config.NewDefaultNamespaceConfig("test.namespace")
ns.DynamicBucketTemplate = config.NewDefaultBucketConfig(config.DynamicBucketTemplateName)
ns.DynamicBucketTemplate.Size = 100000000000
ns.DynamicBucketTemplate.FillRate = 100000000
b := config.NewDefaultBucketConfig("xyz")
config.AddBucket(ns, b)
config.AddNamespace(cfg, ns)
ns = config.NewDefaultNamespaceConfig("test.namespace2")
ns.DefaultBucket = config.NewDefaultBucketConfig(config.DefaultBucketName)
b = config.NewDefaultBucketConfig("xyz")
config.AddBucket(ns, b)
config.AddNamespace(cfg, ns)
server := quotaservice.New(memory.NewBucketFactory(),
config.NewMemoryConfig(cfg),
grpc.New(GRPC_SERVER))
server.SetStatsListener(stats.NewMemoryStatsListener())
server.Start()
// Serve Admin Console
logging.Printf("Starting admin server on %v\n", ADMIN_SERVER)
sm := http.NewServeMux()
server.ServeAdminConsole(sm, "admin/public", true)
go func() {
http.ListenAndServe(ADMIN_SERVER, sm)
}()
// Block until SIGTERM, SIGKILL or SIGINT
sigs := make(chan os.Signal, 1)
signal.Notify(sigs, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGINT)
var shutdown sync.WaitGroup
shutdown.Add(1)
go func() {
<-sigs
shutdown.Done()
}()
shutdown.Wait()
server.Stop()
}
// createNewNamedBucket creates a new, named bucket. May return nil if the named bucket is dynamic,
// and the namespace has already reached its maxDynamicBuckets setting.
func (bc *bucketContainer) createNewNamedBucket(namespace, bucketName string, ns *namespace) *expirableBucket {
bCfg := ns.cfg.Buckets[bucketName]
dyn := false
if bCfg == nil {
// Dynamic.
numDynamicBuckets := bc.countDynamicBuckets(namespace)
if numDynamicBuckets >= ns.cfg.MaxDynamicBuckets && ns.cfg.MaxDynamicBuckets > 0 {
logging.Printf("Bucket %v:%v numDynamicBuckets=%v maxDynamicBuckets=%v. Not creating more dynamic buckets.",
namespace, bucketName, numDynamicBuckets, ns.cfg.MaxDynamicBuckets)
return nil
}
dyn = true
bCfg = ns.cfg.DynamicBucketTemplate
}
return bc.createNewNamedBucketFromCfg(namespace, bucketName, ns, bCfg, dyn)
}
func (z *ZkConfigPersister) waitForEvents(watch *ZkWatch) {
defer z.wg.Done()
for {
select {
case event := <-watch.channel:
if event.Err != nil {
logging.Print("Received error from zookeeper", event)
}
case <-watch.stopper:
return
}
channel, err := watch.listener()
if err != nil {
logging.Printf("Received error from zookeeper executing listener: %+v", err)
continue
}
watch.channel = channel
}
}
// ServeAdminConsole serves up an admin console for an Administrable using Go's built-in HTTP server
// library. `assetsDirectory` contains HTML templates and other UI assets. If empty, no UI will be
// served, and only REST endpoints under `/api/` will be served.
func ServeAdminConsole(a Administrable, mux *http.ServeMux, assetsDirectory string, development bool) {
if assetsDirectory != "" {
msg := "Serving assets from %s"
if development {
msg += " (in development mode)"
}
logging.Printf(msg, assetsDirectory)
mux.Handle("/", loggingHandler(http.RedirectHandler("/admin/", 301)))
mux.Handle("/admin/", loggingHandler(NewUIHandler(a, assetsDirectory, development)))
mux.Handle("/js/", loggingHandler(http.FileServer(http.Dir(assetsDirectory))))
mux.Handle("/favicon.ico", http.NotFoundHandler())
} else {
logging.Print("Not serving admin web UI.")
mux.Handle("/", loggingHandler(http.NotFoundHandler()))
}
bucketsHandler := NewBucketsAPIHandler(a)
namespacesHandler := NewNamespacesAPIHandler(a)
apiHandler := loggingHandler(jsonResponseHandler(apiRequestHandler(
namespacesHandler, bucketsHandler)))
mux.Handle("/api", apiHandler)
mux.Handle("/api/", apiHandler)
statsHandler := loggingHandler(jsonResponseHandler(NewStatsAPIHandler(a)))
mux.Handle("/api/stats", statsHandler)
mux.Handle("/api/stats/", statsHandler)
configsHandler := loggingHandler(jsonResponseHandler(NewConfigsAPIHandler(a)))
mux.Handle("/api/configs", configsHandler)
mux.Handle("/api/configs/", configsHandler)
}
// Tries to create the configuration path, if it doesn't exist
// It tries multiple times in case there's a race with another quotaservice node coming up
func createPath(conn *zk.Conn, path string) (err error) {
for i := 0; i < createRetries; i++ {
exists, _, err := conn.Exists(path)
if exists && err == nil {
return nil
}
_, err = conn.Create(path, []byte{}, 0, zk.WorldACL(zk.PermAll))
if err == nil {
return nil
}
logging.Printf("Could not create zk path, sleeping for 100ms")
time.Sleep(100 * time.Millisecond)
}
if err == nil {
err = errors.New("could not create and get path " + path)
}
return err
}
func writeJSONOk(w http.ResponseWriter) {
if _, e := w.Write(emptyJsonResponse); e != nil {
logging.Printf("Error writing JSON! %+v", e)
}
}