def LeaseFlowForProcessing(self, client_id, flow_id, processing_time): """Marks a flow as being processed on this worker and returns it.""" rdf_flow = self.ReadFlowObject(client_id, flow_id) if rdf_flow.parent_hunt_id: rdf_hunt = self.ReadHuntObject(rdf_flow.parent_hunt_id) if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing( rdf_hunt.hunt_state): raise db.ParentHuntIsNotRunningError(client_id, flow_id, rdf_hunt.hunt_id, rdf_hunt.hunt_state) now = rdfvalue.RDFDatetime.Now() if rdf_flow.processing_on and rdf_flow.processing_deadline > now: raise ValueError("Flow %s on client %s is already being processed." % (flow_id, client_id)) processing_deadline = now + processing_time process_id_string = utils.ProcessIdString() self.UpdateFlow( client_id, flow_id, processing_on=process_id_string, processing_since=now, processing_deadline=processing_deadline) rdf_flow.processing_on = process_id_string rdf_flow.processing_since = now rdf_flow.processing_deadline = processing_deadline return rdf_flow
def ReadFlowForProcessing(self, client_id, flow_id, processing_time): """Marks a flow as being processed on this worker and returns it.""" rdf_flow = self.ReadFlowObject(client_id, flow_id) # TODO(user): remove the check for a legacy hunt prefix as soon as # AFF4 is gone. if rdf_flow.parent_hunt_id and not rdf_flow.parent_hunt_id.startswith("H:"): rdf_hunt = self.ReadHuntObject(rdf_flow.parent_hunt_id) if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing( rdf_hunt.hunt_state): raise db.ParentHuntIsNotRunningError( client_id, flow_id, rdf_hunt.hunt_id, rdf_hunt.hunt_state) now = rdfvalue.RDFDatetime.Now() if rdf_flow.processing_on and rdf_flow.processing_deadline > now: raise ValueError("Flow %s on client %s is already being processed." % (client_id, flow_id)) processing_deadline = now + processing_time process_id_string = utils.ProcessIdString() self.UpdateFlow( client_id, flow_id, processing_on=process_id_string, processing_since=now, processing_deadline=processing_deadline) rdf_flow.processing_on = process_id_string rdf_flow.processing_since = now rdf_flow.processing_deadline = processing_deadline return rdf_flow
def StartHuntFlowOnClient(client_id, hunt_id): """Starts a flow corresponding to a given hunt on a given client.""" hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) hunt_obj = CompleteHuntIfExpirationTimeReached(hunt_obj) # There may be a little race between foreman rules being removed and # foreman scheduling a client on an (already) paused hunt. Making sure # we don't lose clients in such a race by accepting clients for paused # hunts. if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing( hunt_obj.hunt_state): return if hunt_obj.args.hunt_type == hunt_obj.args.HuntType.STANDARD: hunt_args = hunt_obj.args.standard if hunt_obj.client_rate > 0: # Given that we use caching in _GetNumClients and hunt_obj may be updated # in another process, we have to account for cases where num_clients_diff # may go below 0. num_clients_diff = max( 0, _GetNumClients(hunt_obj.hunt_id) - hunt_obj.num_clients_at_start_time) next_client_due_msecs = int(num_clients_diff / hunt_obj.client_rate * 60e6) start_at = rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch( hunt_obj.last_start_time.AsMicrosecondsSinceEpoch() + next_client_due_msecs) else: start_at = None # TODO(user): remove client_rate support when AFF4 is gone. # In REL_DB always work as if client rate is 0. flow_cls = registry.FlowRegistry.FlowClassByName(hunt_args.flow_name) flow_args = hunt_args.flow_args if hunt_args.HasField( "flow_args") else None flow.StartFlow( client_id=client_id, creator=hunt_obj.creator, cpu_limit=hunt_obj.per_client_cpu_limit, network_bytes_limit=hunt_obj.per_client_network_bytes_limit, flow_cls=flow_cls, flow_args=flow_args, start_at=start_at, parent_hunt_id=hunt_id) if hunt_obj.client_limit: if _GetNumClients(hunt_obj.hunt_id) >= hunt_obj.client_limit: PauseHunt(hunt_id) elif hunt_obj.args.hunt_type == hunt_obj.args.HuntType.VARIABLE: raise NotImplementedError() else: raise UnknownHuntTypeError("Can't determine hunt type when starting " "hunt %s on client %s." % (client_id, hunt_id))