@@ -114,32 +114,85 @@ def _make_ctx_options(ctx_options, config_cls=ContextOptions):
114114
115115
116116class AutoBatcher (object ):
117+ """Batches multiple async calls if they share the same rpc options.
118+
119+ Here is an example to explain what this class does.
120+
121+ Life of a key.get_async(options) API call:
122+ *) Key gets the singleton Context instance and invokes Context.get.
123+ *) Context.get calls Context._get_batcher.add(key, options). This
124+ returns a future "fut" as the return value of key.get_async.
125+ At this moment, key.get_async returns.
126+
127+ *) When more than "limit" number of _get_batcher.add() was called,
128+ _get_batcher invokes its self._todo_tasklet, Context._get_tasklet,
129+ with the list of keys seen so far.
130+ *) Context._get_tasklet fires a MultiRPC and waits on it.
131+ *) Upon MultiRPC completion, Context._get_tasklet passes on the results
132+ to the respective "fut" from key.get_async.
133+
134+ *) If user calls "fut".get_result() before "limit" number of add() was called,
135+ "fut".get_result() will repeatedly call eventloop.run1().
136+ *) After processing immediate callbacks, eventloop will run idlers.
137+ AutoBatcher._on_idle is an idler.
138+ *) _on_idle will run the "todo_tasklet" before the batch is full.
139+
140+ So the engine is todo_tasklet, which is a proxy tasklet that can combine
141+ arguments into batches and passes along results back to respective futures.
142+ This class is mainly a helper that invokes todo_tasklet with the right
143+ arguments at the right time.
144+ """
117145
118146 def __init__ (self , todo_tasklet , limit ):
119- # todo_tasklet is a tasklet to be called with list of (future, arg) pairs
147+ """Init.
148+
149+ Args:
150+ todo_tasklet: the tasklet that actually fires RPC and waits on a MultiRPC.
151+ It should take a list of (future, arg) pairs and an "options" as
152+ arguments. "options" are rpc options.
153+ limit: max number of items to batch for each distinct value of "options".
154+ """
120155 self ._todo_tasklet = todo_tasklet
121- self ._limit = limit # No more than this many per callback
122- self ._queues = {} # Map options to lists of (future, arg) tuples
123- self ._running = [] # Currently running tasklets
124- self ._cache = {} # Cache of in-flight todo_tasklet futures
156+ self ._limit = limit
157+ # A map from "options" to a list of (future, arg) tuple.
158+ # future is the future return from a single async operations.
159+ self ._queues = {}
160+ self ._running = [] # A list of in-flight todo_tasklet futures.
161+ self ._cache = {} # Cache of in-flight todo_tasklet futures.
125162
126163 def __repr__ (self ):
127164 return '%s(%s)' % (self .__class__ .__name__ , self ._todo_tasklet .__name__ )
128165
129166 def run_queue (self , options , todo ):
167+ """Actually run the _todo_tasklet."""
130168 utils .logging_debug ('AutoBatcher(%s): %d items' ,
131169 self ._todo_tasklet .__name__ , len (todo ))
132- fut = self ._todo_tasklet (todo , options )
133- self ._running .append (fut )
170+ batch_fut = self ._todo_tasklet (todo , options )
171+ self ._running .append (batch_fut )
134172 # Add a callback when we're done.
135- fut .add_callback (self ._finished_callback , fut , todo )
173+ batch_fut .add_callback (self ._finished_callback , batch_fut , todo )
136174
137175 def _on_idle (self ):
176+ """An idler eventloop can run.
177+
178+ Eventloop calls this when it has finished processing all immediate
179+ callbacks. This method runs _todo_tasklet even before the batch is full.
180+ """
138181 if not self .action ():
139182 return None
140183 return True
141184
142185 def add (self , arg , options = None ):
186+ """Adds an arg and gets back a future.
187+
188+ Args:
189+ arg: one argument for _todo_tasklet.
190+ options: rpc options.
191+
192+ Return:
193+ An instance of future, representing the result of running
194+ _todo_tasklet without batching.
195+ """
143196 fut = tasklets .Future ('%s.add(%s, %s)' % (self , arg , options ))
144197 todo = self ._queues .get (options )
145198 if todo is None :
@@ -171,14 +224,24 @@ def action(self):
171224 self .run_queue (options , todo )
172225 return True
173226
174- def _finished_callback (self , fut , todo ):
175- self ._running .remove (fut )
176- err = fut .get_exception ()
227+ def _finished_callback (self , batch_fut , todo ):
228+ """Passes exception along.
229+
230+ Args:
231+ batch_fut: the batch future returned by running todo_tasklet.
232+ todo: (fut, option) pair. fut is the future return by each add() call.
233+
234+ If the batch fut was successful, it has already called fut.set_result()
235+ on other individual futs. This method only handles when the batch fut
236+ encountered an exception.
237+ """
238+ self ._running .remove (batch_fut )
239+ err = batch_fut .get_exception ()
177240 if err is not None :
178- tb = fut .get_traceback ()
179- for (f , _ ) in todo :
180- if not f .done ():
181- f .set_exception (err , tb )
241+ tb = batch_fut .get_traceback ()
242+ for (fut , _ ) in todo :
243+ if not fut .done ():
244+ fut .set_exception (err , tb )
182245
183246 @tasklets .tasklet
184247 def flush (self ):
@@ -647,9 +710,8 @@ def get(self, key, **ctx_options):
647710 if use_cache :
648711 self ._load_from_cache_if_available (key )
649712 if mvalue not in _LOCKED , None ):
650- cls = model .Model ._kind_map .get (key .kind ())
651- if cls is None :
652- raise TypeError ('Cannot find model class for kind %s' % key .kind ())
713+ cls = model .Model ._lookup_model (key .kind (),
714+ self ._conn .adapter .default_model )
653715 pb = entity_pb .EntityProto ()
654716
655717 try :
0 commit comments