1# Copyright 2016 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Feedback query delegate interfaces and implementation registry."""
6
7import multiprocessing
8
9import common
10from autotest_lib.client.common_lib.feedback import client
11
12
13# Mapping of query identifiers to delegate classes.
14_query_delegate_registry = {}
15
16
17class _QueryDelegate(object):
18    """A base class for query delegates."""
19
20    _query_count = multiprocessing.Value('d', 0)
21
22    def __init__(self, test, dut, multiplexer, atomic=True):
23        """Constructs the delegate.
24
25        @param test: The name of the test.
26        @param dut: The name of the DUT.
27        @param multiplexer: Feedback request multiplexer object.
28        @param atomic: Whether this is an atomic query.
29        """
30        super(_QueryDelegate, self).__init__()
31        self.test = test
32        self.dut = dut
33        self._multiplexer = multiplexer
34        self._atomic = atomic
35
36        # Assign a unique query number.
37        with self._query_count.get_lock():
38            self._query_num = self._query_count.value
39            self._query_count.value += 1
40
41
42    def _process_request(self, req):
43        """Submits a given request to the multiplexer for processing."""
44        return self._multiplexer.process_request(req, self._query_num,
45                                                 self._atomic)
46
47
48    def prepare(self, **kwargs):
49        """Delegate for a query's prepare() method."""
50        return self._prepare_impl(**kwargs)
51
52
53    def _prepare_impl(self, **kwargs):
54        """Concrete implementation of the query's prepare() call."""
55        raise NotImplementedError
56
57
58    def validate(self, **kwargs):
59        """Delegate for a query's validate() method.
60
61        This clears the atomic sequence with the multiplexer to make sure it
62        isn't blocked waiting for more requests from this query.
63        """
64        try:
65            return self._validate_impl(**kwargs)
66        finally:
67            if self._atomic:
68                self._multiplexer.end_atomic_seq(self._query_num)
69
70
71    def _validate_impl(self, **kwargs):
72        """Concrete implementation of the query's validate() call."""
73        raise NotImplementedError
74
75
76class OutputQueryDelegate(_QueryDelegate):
77    """A base class for output query delegates."""
78
79
80class InputQueryDelegate(_QueryDelegate):
81    """A base class for input query delegates."""
82
83    def emit(self):
84        """Delegate for an input query's emit() method."""
85        return self._emit_impl()
86
87
88    def _emit_impl(self):
89        """Concrete implementation of the query's emit() call."""
90        raise NotImplementedError
91
92
93def register_delegate_cls(query_id, delegate_cls):
94    """Registers a delegate class with a given query identifier.
95
96    @param query_id: Query identifier constant.
97    @param delegate_cls: The class implementing a delegate for this query.
98    """
99    _query_delegate_registry[query_id] = delegate_cls
100
101
102def get_delegate_cls(query_id):
103    """Returns a query delegate class for a given query type.
104
105    @param query_id: A query type identifier.
106
107    @return A query delegate class.
108
109    @raise ValueError: Unknown query type.
110    @raise NotImplementedError: Query type not supported.
111    """
112    if query_id not in client.ALL_QUERIES:
113        raise ValueError
114    if query_id not in _query_delegate_registry:
115        raise NotImplementedError
116    return _query_delegate_registry[query_id]
117