backends/counter/counter_v1_with_shutdown.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright 2011 Google Inc. All Rights Reserved.
   4 #
   5 # Licensed under the Apache License, Version 2.0 (the "License");
   6 # you may not use this file except in compliance with the License.
   7 # You may obtain a copy of the License at
   8 #
   9 #     http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing, software
  12 # distributed under the License is distributed on an "AS IS" BASIS,
  13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 # See the License for the specific language governing permissions and
  15 # limitations under the License.
  16 # vim: set ts=4 sw=4 et tw=79:
  17
  18 """Implementation of a counter with persistent storage."""
  19
  20 __author__ = 'Greg Darke <darke@google.com>'
  21
  22 import logging
  23 import math
  24 import time
  25
  26 from google.appengine.api import backends
  27 from google.appengine.api import runtime
  28 from google.appengine.api import taskqueue
  29 from google.appengine.ext import db
  30 from google.appengine.ext import webapp
  31 from google.appengine.ext.webapp.util import run_wsgi_app
  32
  33
  34 class CounterModel(db.Model):
  35     value = db.IntegerProperty(default=0)
  36     _dirty = False
  37
  38     @classmethod
  39     def get_or_new(cls, name):
  40         model = cls.get_by_key_name(name)
  41         if model is None:
  42             model = cls(key_name=name)
  43         return model
  44
  45
  46 class CounterStore(object):
  47     def __init__(self):
  48         self._store = {}
  49         self._has_shutdown = False
  50         self._batch_size = 100
  51
  52     def get_value(self, name):
  53         if name not in self._store:
  54             model = CounterModel.get_or_new(name)
  55             self._store[name] = model
  56         else:
  57             model = self._store[name]
  58         return model
  59
  60     def inc_value(self, name, delta):
  61         model = self.get_value(name)
  62         model.value += delta
  63         model._dirty = True
  64         if self._has_shutdown:
  65             # Since the shutdown hook may be called at any time, we need to
  66             # protect ourselves against this.
  67             model.put()
  68         return model
  69
  70     def _put_counters(self, counters):
  71         db.put(counters)
  72         for counter in counters:
  73             counter._dirty = False
  74
  75     def _write_in_batches(self, counters):
  76         """Write out the dirty entries from 'counters' in batches.
  77
  78         The batch size is determined by self._batch_size.
  79
  80         Args:
  81           counters: An iterable containing instances of CounterModel.
  82         """
  83         to_put = []
  84         for counter in counters:
  85             if counter._dirty:
  86                 to_put.append(counter)
  87
  88                 if len(to_put) >= self._batch_size:
  89                     self._put_counters(to_put)
  90                     to_put = []
  91
  92         if to_put:
  93             self._put_counters(to_put)
  94
  95     def flush_to_datastore(self):
  96         """Write the dirty entries from _store to datastore."""
  97         self._write_in_batches(self._store.itervalues())
  98         self._dirty = False
  99
 100     def shutdown_hook(self):
 101         """Ensures all counters are written to datastore."""
 102         self.flush_to_datastore()
 103         self._has_shutdown = True
 104
 105
 106 class StartHandler(webapp.RequestHandler):
 107     """Handler for '/_ah/start'.
 108
 109     This url is called once when the backend is started.
 110     """
 111     def get(self):
 112         runtime.set_shutdown_hook(_counter_store.shutdown_hook)
 113
 114
 115 class CounterHandler(webapp.RequestHandler):
 116     """Handler for counter/{get,inc,dec}.
 117
 118     This handler is protected by login: admin in app.yaml.
 119     """
 120
 121     def _write_error(self, error_message):
 122         self.response.error(400)
 123         self.response.out.write(error_message)
 124
 125     def post(self, method):
 126         """Handler a post to counter/{get,inc,dec}.
 127
 128         The 'method' parameter is parsed from the url by a regex capture group.
 129
 130         Args:
 131           method: The type of operation to perform against the counter store.
 132               It must be one of 'get', 'inc' or 'dec'.
 133         """
 134         key_name = self.request.get('name')
 135         delta = self.request.get('delta')
 136
 137         if not key_name:
 138             self._write_error('Request did not have a "key_name" parameter.')
 139             return
 140
 141         if method == 'get':
 142             model = _counter_store.get_value(key_name)
 143         else:
 144             if not delta:
 145                 self._write_error('Request did not have a "delta" parameter.')
 146                 return
 147             try:
 148                 delta = int(delta)
 149             except ValueError:
 150                 self._write_error('"delta" is not an integer.')
 151                 return
 152
 153             if method == 'inc':
 154                 model = _counter_store.inc_value(key_name, delta)
 155             elif method == 'dec':
 156                 model = _counter_store.inc_value(key_name, -delta)
 157
 158         self.response.headers['Content-Type'] = 'text/plain'
 159         self.response.out.write('%d' % model.value)
 160
 161
 162 _handlers = [(r'/_ah/start', StartHandler),
 163              (r'/backend/counter/(get|inc|dec)$', CounterHandler)]
 164
 165 application = webapp.WSGIApplication(_handlers)
 166 _counter_store = CounterStore()
 167
 168
 169 def main():
 170     run_wsgi_app(application)
 171
 172 if __name__ == '__main__':
 173     main()