RPC/test_delay: Use callable for timeout calculation

This avoids having to override the function in the RPC runner.

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: René Nussbaumer <rn@google.com>

autotools/build-rpc

@@ -129,9 +129,6 @@ def _WriteBaseClass(sw, clsname, calls):
 
       funcargs.extend(map(compat.fst, args))
 
-      assert "read_timeout" not in funcargs
-      funcargs.append("read_timeout=%s" % timeout)
-
       funcargs.append("_def=_CALLS[%r]" % name)
 
       funcdef = "def call_%s(%s):" % (name, utils.CommaJoin(funcargs))
@@ -154,7 +151,7 @@ def _WriteBaseClass(sw, clsname, calls):
         else:
           buf.write("node_list")
 
-        buf.write(", read_timeout, [%s])" %
+        buf.write(", [%s])" %
                   # Function arguments
                   utils.CommaJoin(map(compat.fst, args)))
 

lib/rpc.py

@@ -430,17 +430,22 @@ class _RpcClientBase:
     else:
       return encoder_fn(argkind)(value)
 
-  def _Call(self, cdef, node_list, timeout, args):
+  def _Call(self, cdef, node_list, args):
     """Entry point for automatically generated RPC wrappers.
 
     """
-    (procedure, _, _, argdefs, postproc_fn, _) = cdef
+    (procedure, _, timeout, argdefs, postproc_fn, _) = cdef
+
+    if callable(timeout):
+      read_timeout = timeout(args)
+    else:
+      read_timeout = timeout
 
     body = serializer.DumpJson(map(self._encoder,
                                    zip(map(compat.snd, argdefs), args)),
                                indent=False)
 
-    result = self._proc(node_list, procedure, body, read_timeout=timeout)
+    result = self._proc(node_list, procedure, body, read_timeout=read_timeout)
 
     if postproc_fn:
       return postproc_fn(result)
@@ -620,20 +625,6 @@ class RpcRunner(_RpcClientBase,
     """
     return self._InstDict(instance, osp=osparams)
 
-  #
-  # Begin RPC calls
-  #
-
-  def call_test_delay(self, node_list, duration): # pylint: disable=W0221
-    """Sleep for a fixed time on given node(s).
-
-    This is a multi-node call.
-
-    """
-    # TODO: Use callable timeout calculation
-    return _generated_rpc.RpcClientDefault.call_test_delay(self,
-      node_list, duration, read_timeout=int(duration + 5))
-
 
 class JobQueueRunner(_RpcClientBase, _generated_rpc.RpcClientJobQueue):
   """RPC wrappers for job queue.

lib/rpc_defs.py

@@ -24,7 +24,8 @@ RPC definition fields:
 
   - Name as string
   - L{SINGLE} for single-node calls, L{MULTI} for multi-node
-  - Timeout (e.g. L{TMO_NORMAL})
+  - Timeout (e.g. L{TMO_NORMAL}), or callback receiving all arguments in a
+    tuple to calculate timeout
   - List of arguments as tuples
 
     - Name as string
@@ -149,6 +150,13 @@ def _ImpExpStatusPostProc(result):
   return result
 
 
+def _TestDelayTimeout((duration, )):
+  """Calculate timeout for "test_delay" RPC.
+
+  """
+  return int(duration + 5)
+
+
 _FILE_STORAGE_CALLS = [
   ("file_storage_dir_create", SINGLE, TMO_FAST, [
     ("file_storage_dir", None, "File storage directory"),
@@ -457,7 +465,7 @@ _MISC_CALLS = [
     ], None, "Call an iallocator on a remote node"),
   ("test_delay", MULTI, None, [
     ("duration", None, None),
-    ], None, "Sleep for a fixed time on given node(s)"),
+    ], _TestDelayTimeout, "Sleep for a fixed time on given node(s)"),
   ("hypervisor_validate_params", MULTI, TMO_NORMAL, [
     ("hvname", None, "Hypervisor name"),
     ("hvfull", None, "Parameters to be validated"),