Locust - core.py模块源码剖析

core.py模块模块是整个Locust的核心部分代码，包含了HttpLocust类和TaskSet类等。

目录

task函数

Locust类

基于Locust类的HttpLocust类

TaskSetMeta元类

TaskSet类

 

task函数

def task(weight=1):
    """
    常被用作装饰器，用于声明是TaskSet类中的任务，并分配所占比例
    class ForumPage(TaskSet):
        @task(100)
        def read_thread(self):
            pass
            
        @task(7)
        def create_thread(self):
            pass
    """
    def decorator_func(func):
        func.locust_task_weight = weight
        return func
        
    if callable(weight):
        func = weight
        weight = 1
        return decorator_func(func)
    else:
        return decorator_func

Locust类

class Locust(object):
    """
    每个Locust类表示一个对系统施压的用户
    每个用户的行为需要在类内定义task_set类
    """
    
    host = None
    min_wait = 1000
    max_wait = 1000
    task_set = None
    stop_timeout = None
    weight = 10
        
    client = NoClientWarningRaiser()
    _catch_exceptions = True
    
    def __init__(self):
        super(Locust, self).__init__()
    
    def run(self):
        try:
            self.task_set(self).run()
        except StopLocust:
            pass
        except (RescheduleTask, RescheduleTaskImmediately) as e:
            six.reraise(LocustError, LocustError("A task inside a Locust class' main TaskSet (`%s.task_set` of type `%s`) seems to have called interrupt() or raised an InterruptTaskSet exception. The interrupt() function is used to hand over execution to a parent TaskSet, and should never be called in the main TaskSet which a Locust class' task_set attribute points to." % (type(self).__name__, self.task_set.__name__)), sys.exc_info()[2])

基于Locust类的HttpLocust类

class HttpLocust(Locust):
    """
    在Locust下继承得到的，包含一个client属性，用于发送HTTP请求
    """
    
    client = None
    
    def __init__(self):
        super(HttpLocust, self).__init__()
        if self.host is None:
            raise LocustError("You must specify the base host. Either in the host attribute in the Locust class, or on the command line using the --host option.")
        
        self.client = HttpSession(base_url=self.host)

TaskSetMeta元类

class TaskSetMeta(type):
    """
    Locust的核心元类，主要功能是可以让Locust类来指定特定的执行比例
    """
    def __new__(mcs, classname, bases, classDict):
        new_tasks = []
        # 依次获取task
        for base in bases:
            if hasattr(base, "tasks") and base.tasks:
                new_tasks += base.tasks
        
        if "tasks" in classDict and classDict["tasks"] is not None:
            tasks = classDict["tasks"]
            if isinstance(tasks, dict):
                tasks = six.iteritems(tasks)
            
            for task in tasks:
                if isinstance(task, tuple):
                    task, count = task
                    for i in xrange(0, count):
                        new_tasks.append(task)
                else:
                    new_tasks.append(task)
        
        for item in six.itervalues(classDict):
            if hasattr(item, "locust_task_weight"):
                for i in xrange(0, item.locust_task_weight):
                    new_tasks.append(item)
        # 获取所有按照比例分配好的new_tasks
        classDict["tasks"] = new_tasks
        
        return type.__new__(mcs, classname, bases, classDict)

TaskSet类

@six.add_metaclass(TaskSetMeta)
class TaskSet(object):
    """
    TaskSet类主要是用于定义一组提供Locust User执行的操作。
    """
    
    tasks = []
    min_wait = None
    max_wait = None
    locust = None
    parent = None
 
    def __init__(self, parent):
        self._task_queue = []
        self._time_start = time()
        
        if isinstance(parent, TaskSet):
            self.locust = parent.locust
        elif isinstance(parent, Locust):
            self.locust = parent
        else:
            raise LocustError("TaskSet should be called with Locust instance or TaskSet instance as first argument")
 
        self.parent = parent
        
        if not self.min_wait:
            self.min_wait = self.locust.min_wait
        if not self.max_wait:
            self.max_wait = self.locust.max_wait
 
    def run(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        
        try:
            if hasattr(self, "on_start"):
                self.on_start()
        except InterruptTaskSet as e:
            if e.reschedule:
                six.reraise(RescheduleTaskImmediately, RescheduleTaskImmediately(e.reschedule), sys.exc_info()[2])
            else:
                six.reraise(RescheduleTask, RescheduleTask(e.reschedule), sys.exc_info()[2])
        
        while (True):
            try:
                if self.locust.stop_timeout is not None and time() - self._time_start > self.locust.stop_timeout:
                    return
        
                if not self._task_queue:
                    self.schedule_task(self.get_next_task())
                
                try:
                    self.execute_next_task()
                except RescheduleTaskImmediately:
                    pass
                except RescheduleTask:
                    self.wait()
                else:
                    self.wait()
            except InterruptTaskSet as e:
                if e.reschedule:
                    six.reraise(RescheduleTaskImmediately, RescheduleTaskImmediately(e.reschedule), sys.exc_info()[2])
                else:
                    six.reraise(RescheduleTask, RescheduleTask(e.reschedule), sys.exc_info()[2])
            except StopLocust:
                raise
            except GreenletExit:
                raise
            except Exception as e:
                events.locust_error.fire(locust_instance=self, exception=e, tb=sys.exc_info()[2])
                if self.locust._catch_exceptions:
                    sys.stderr.write("\n" + traceback.format_exc())
                    self.wait()
                else:
                    raise
    
    def execute_next_task(self):
        task = self._task_queue.pop(0)
        self.execute_task(task["callable"], *task["args"], **task["kwargs"])
    
    def execute_task(self, task, *args, **kwargs):
        if hasattr(task, "__self__") and task.__self__ == self:
            task(*args, **kwargs)
        elif hasattr(task, "tasks") and issubclass(task, TaskSet):
            task(self).run(*args, **kwargs)
        else:
            task(self, *args, **kwargs)
    
    def schedule_task(self, task_callable, args=None, kwargs=None, first=False):
        task = {"callable":task_callable, "args":args or [], "kwargs":kwargs or {}}
        if first:
            self._task_queue.insert(0, task)
        else:
            self._task_queue.append(task)
    
    def get_next_task(self):
        return random.choice(self.tasks)
    
    def wait(self):
        millis = random.randint(self.min_wait, self.max_wait)
        seconds = millis / 1000.0
        self._sleep(seconds)
 
    def _sleep(self, seconds):
        gevent.sleep(seconds)
    
    def interrupt(self, reschedule=True):
        raise InterruptTaskSet(reschedule)
    
    @property
    def client(self):
        return self.locust.client