nodes: 0, module; 1, function_definition; 2, function_name:DrainTaskSchedulerQueueForClient; 3, parameters; 4, block; 5, identifier:self; 6, identifier:client; 7, default_parameter; 8, expression_statement; 9, if_statement; 10, if_statement; 11, expression_statement; 12, expression_statement; 13, comment:# Drain the queue for this client; 14, if_statement; 15, expression_statement; 16, if_statement; 17, return_statement; 18, identifier:max_count; 19, None; 20, comment:"""Drains the client's Task Scheduler queue.
  
      1) Get all messages in the client queue.
      2) Sort these into a set of session_ids.
      3) Use data_store.DB.ResolvePrefix() to query all requests.
      4) Delete all responses for retransmitted messages (if needed).
  
      Args:
         client: The ClientURN object specifying this client.
         max_count: The maximum number of messages we will issue for the client.
           If not given, uses self.max_queue_size .
  
      Returns:
         The tasks respresenting the messages returned. If we can not send them,
         we can reschedule them for later.
      """; 21, comparison_operator:max_count is None; 22, block; 23, comparison_operator:max_count <= 0; 24, block; 25, assignment; 26, assignment; 27, call; 28, block; 29, else_clause; 30, call; 31, identifier:result; 32, block; 33, identifier:result; 34, identifier:max_count; 35, None; 36, expression_statement; 37, identifier:max_count; 38, integer:0; 39, return_statement; 40, identifier:client; 41, call; 42, identifier:start_time; 43, call; 44, attribute; 45, argument_list; 46, expression_statement; 47, expression_statement; 48, block; 49, attribute; 50, argument_list; 51, expression_statement; 52, assignment; 53, list; 54, attribute; 55, argument_list; 56, attribute; 57, argument_list; 58, identifier:data_store; 59, identifier:RelationalDBEnabled; 60, assignment; 61, assignment; 62, expression_statement; 63, expression_statement; 64, expression_statement; 65, expression_statement; 66, for_statement; 67, if_statement; 68, call; 69, identifier:IncrementCounter; 70, string:"grr_messages_sent"; 71, call; 72, call; 73, identifier:max_count; 74, attribute; 75, identifier:rdf_client; 76, identifier:ClientURN; 77, identifier:client; 78, identifier:time; 79, identifier:time; 80, identifier:action_requests; 81, call; 82, identifier:result; 83, list_comprehension; 84, assignment; 85, assignment; 86, assignment; 87, assignment; 88, identifier:task; 89, identifier:new_tasks; 90, block; 91, identifier:check_before_sending; 92, block; 93, attribute; 94, argument_list; 95, identifier:len; 96, argument_list; 97, attribute; 98, argument_list; 99, identifier:self; 100, identifier:max_queue_size; 101, attribute; 102, argument_list; 103, call; 104, for_in_clause; 105, identifier:new_tasks; 106, call; 107, identifier:initial_ttl; 108, attribute; 109, identifier:check_before_sending; 110, list; 111, identifier:result; 112, list; 113, if_statement; 114, with_statement; 115, identifier:stats_collector_instance; 116, identifier:Get; 117, identifier:result; 118, identifier:logging; 119, identifier:debug; 120, string:"Drained %d messages for %s in %s seconds."; 121, call; 122, identifier:client; 123, binary_operator:time.time() - start_time; 124, attribute; 125, identifier:LeaseClientActionRequests; 126, call; 127, keyword_argument; 128, keyword_argument; 129, attribute; 130, argument_list; 131, identifier:r; 132, identifier:action_requests; 133, attribute; 134, argument_list; 135, call; 136, identifier:task_ttl; 137, comparison_operator:task.task_ttl < initial_ttl - 1; 138, comment:# This message has been leased before.; 139, block; 140, else_clause; 141, with_clause; 142, block; 143, identifier:len; 144, argument_list; 145, call; 146, identifier:start_time; 147, identifier:data_store; 148, identifier:REL_DB; 149, attribute; 150, argument_list; 151, identifier:lease_time; 152, call; 153, identifier:limit; 154, identifier:max_count; 155, identifier:rdf_flow_objects; 156, identifier:GRRMessageFromClientActionRequest; 157, identifier:r; 158, call; 159, identifier:QueryAndOwn; 160, keyword_argument; 161, keyword_argument; 162, keyword_argument; 163, attribute; 164, argument_list; 165, attribute; 166, binary_operator:initial_ttl - 1; 167, expression_statement; 168, block; 169, with_item; 170, expression_statement; 171, comment:# All messages that don't have a status yet should be sent again.; 172, for_statement; 173, identifier:result; 174, attribute; 175, argument_list; 176, identifier:client; 177, identifier:Basename; 178, attribute; 179, argument_list; 180, attribute; 181, argument_list; 182, identifier:queue; 183, call; 184, identifier:limit; 185, identifier:max_count; 186, identifier:lease_seconds; 187, attribute; 188, identifier:rdf_flows; 189, identifier:GrrMessage; 190, identifier:task; 191, identifier:task_ttl; 192, identifier:initial_ttl; 193, integer:1; 194, call; 195, expression_statement; 196, as_pattern; 197, assignment; 198, identifier:task; 199, identifier:check_before_sending; 200, block; 201, identifier:time; 202, identifier:time; 203, attribute; 204, identifier:FromSeconds; 205, attribute; 206, identifier:queue_manager; 207, identifier:QueueManager; 208, keyword_argument; 209, attribute; 210, argument_list; 211, identifier:self; 212, identifier:message_expiry_time; 213, attribute; 214, argument_list; 215, call; 216, call; 217, as_pattern_target; 218, identifier:status_found; 219, call; 220, if_statement; 221, identifier:rdfvalue; 222, identifier:Duration; 223, identifier:self; 224, identifier:message_expiry_time; 225, identifier:token; 226, attribute; 227, identifier:client; 228, identifier:Queue; 229, identifier:check_before_sending; 230, identifier:append; 231, identifier:task; 232, attribute; 233, argument_list; 234, attribute; 235, argument_list; 236, identifier:manager; 237, attribute; 238, argument_list; 239, comparison_operator:task not in status_found; 240, block; 241, else_clause; 242, identifier:self; 243, identifier:token; 244, identifier:result; 245, identifier:append; 246, identifier:task; 247, identifier:queue_manager; 248, identifier:QueueManager; 249, keyword_argument; 250, identifier:manager; 251, identifier:MultiCheckStatus; 252, identifier:check_before_sending; 253, identifier:task; 254, identifier:status_found; 255, expression_statement; 256, block; 257, identifier:token; 258, attribute; 259, call; 260, expression_statement; 261, identifier:self; 262, identifier:token; 263, attribute; 264, argument_list; 265, call; 266, identifier:result; 267, identifier:append; 268, identifier:task; 269, attribute; 270, argument_list; 271, identifier:manager; 272, identifier:DeQueueClientRequest; 273, identifier:task
code: def DrainTaskSchedulerQueueForClient(self, client, max_count=None):
      """Drains the client's Task Scheduler queue.
  
      1) Get all messages in the client queue.
      2) Sort these into a set of session_ids.
      3) Use data_store.DB.ResolvePrefix() to query all requests.
      4) Delete all responses for retransmitted messages (if needed).
  
      Args:
         client: The ClientURN object specifying this client.
         max_count: The maximum number of messages we will issue for the client.
           If not given, uses self.max_queue_size .
  
      Returns:
         The tasks respresenting the messages returned. If we can not send them,
         we can reschedule them for later.
      """
      if max_count is None:
        max_count = self.max_queue_size
  
      if max_count <= 0:
        return []
  
      client = rdf_client.ClientURN(client)
      start_time = time.time()
      # Drain the queue for this client
      if data_store.RelationalDBEnabled():
        action_requests = data_store.REL_DB.LeaseClientActionRequests(
            client.Basename(),
            lease_time=rdfvalue.Duration.FromSeconds(self.message_expiry_time),
            limit=max_count)
        result = [
            rdf_flow_objects.GRRMessageFromClientActionRequest(r)
            for r in action_requests
        ]
      else:
        new_tasks = queue_manager.QueueManager(token=self.token).QueryAndOwn(
            queue=client.Queue(),
            limit=max_count,
            lease_seconds=self.message_expiry_time)
  
        initial_ttl = rdf_flows.GrrMessage().task_ttl
        check_before_sending = []
        result = []
        for task in new_tasks:
          if task.task_ttl < initial_ttl - 1:
            # This message has been leased before.
            check_before_sending.append(task)
          else:
            result.append(task)
  
        if check_before_sending:
          with queue_manager.QueueManager(token=self.token) as manager:
            status_found = manager.MultiCheckStatus(check_before_sending)
  
            # All messages that don't have a status yet should be sent again.
            for task in check_before_sending:
              if task not in status_found:
                result.append(task)
              else:
                manager.DeQueueClientRequest(task)
  
      stats_collector_instance.Get().IncrementCounter("grr_messages_sent",
                                                      len(result))
      if result:
        logging.debug("Drained %d messages for %s in %s seconds.", len(result),
                      client,
                      time.time() - start_time)
  
      return result