misc.python.materialize.parallel_workload.worker

View Source

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10import random
 11import threading
 12import time
 13from collections import Counter, defaultdict
 14
 15import psycopg
 16import websocket
 17
 18from materialize.data_ingest.query_error import QueryError
 19from materialize.mzcompose.composition import Composition
 20from materialize.parallel_workload.action import (
 21    Action,
 22    ActionList,
 23    ReconnectAction,
 24    ws_connect,
 25)
 26from materialize.parallel_workload.database import Database
 27from materialize.parallel_workload.executor import Executor
 28
 29
 30class Worker:
 31    rng: random.Random
 32    action_list: ActionList | None
 33    actions: list[Action]
 34    weights: list[float]
 35    end_time: float
 36    num_queries: Counter[type[Action]]
 37    num_successes: Counter[type[Action]]
 38    num_skips: Counter[type[Action]]
 39    autocommit: bool
 40    system: bool
 41    exe: Executor | None
 42    ignored_errors: defaultdict[str, Counter[type[Action]]]
 43    composition: Composition | None
 44    occurred_exception: Exception | None
 45
 46    def __init__(
 47        self,
 48        rng: random.Random,
 49        actions: list[Action],
 50        weights: list[float],
 51        end_time: float,
 52        autocommit: bool,
 53        system: bool,
 54        composition: Composition | None,
 55        action_list: ActionList | None = None,
 56    ):
 57        self.rng = rng
 58        self.action_list = action_list
 59        self.actions = actions
 60        self.weights = weights
 61        self.end_time = end_time
 62        self.num_queries = Counter()
 63        # Unlike num_queries, these are never cleared: they feed the
 64        # end-of-run action coverage check.
 65        self.num_successes = Counter()
 66        self.num_skips = Counter()
 67        self.autocommit = autocommit
 68        self.system = system
 69        self.ignored_errors = defaultdict(Counter)
 70        self.composition = composition
 71        self.occurred_exception = None
 72        self.exe = None
 73
 74    def run(
 75        self, host: str, pg_port: int, http_port: int, user: str, database: Database
 76    ) -> None:
 77        # In scenarios with kills and deploys materialized can go down at any
 78        # point during the setup, keep retrying.
 79        for i in range(300):
 80            try:
 81                self.conn = psycopg.connect(
 82                    host=host, port=pg_port, user=user, dbname="materialize"
 83                )
 84                self.conn.autocommit = self.autocommit
 85                cur = self.conn.cursor()
 86                ws = websocket.WebSocket()
 87                ws_conn_id, ws_secret_key = ws_connect(ws, host, http_port, user)
 88                self.exe = Executor(self.rng, cur, ws, database, user=user)
 89                self.exe.set_isolation("SERIALIZABLE")
 90                cur.execute("SET auto_route_catalog_queries TO false")
 91                if self.exe.use_ws:
 92                    self.exe.pg_pid = ws_conn_id
 93                else:
 94                    cur.execute("SELECT pg_backend_pid()")
 95                    self.exe.pg_pid = cur.fetchall()[0][0]
 96            except Exception:
 97                if time.time() > self.end_time:
 98                    return
 99                if i == 299:
100                    raise
101                time.sleep(1)
102            else:
103                break
104        assert self.exe
105
106        while time.time() < self.end_time:
107            action = self.rng.choices(self.actions, self.weights)[0]
108            if not action.applicable(self.exe):
109                continue
110            if self.exe.rollback_next:
111                try:
112                    self.exe.rollback()
113                except QueryError:
114                    # ROLLBACK can itself fail, e.g. cancelled by
115                    # `pg_cancel_backend` or on a broken WS session. Force a
116                    # reconnect rather than leaving a session with an open
117                    # aborted transaction behind.
118                    self.exe.reconnect_next = True
119                    self.exe.rollback_next = False
120                    continue
121                self.exe.rollback_next = False
122            if self.exe.reconnect_next:
123                self.exe.reconnect_next = False
124                # Run as its own action so failures are attributed to
125                # ReconnectAction, not to `action`, which hasn't run yet.
126                self.run_action(
127                    ReconnectAction(self.rng, self.composition, random_role=False)
128                )
129                if self.exe.reconnect_next or self.exe.rollback_next:
130                    # Reconnecting failed with an ignored error. Always retry
131                    # the reconnect, never fall through to the action: the
132                    # old session may hold an aborted transaction that fails
133                    # all statements.
134                    self.exe.reconnect_next = True
135                    self.exe.rollback_next = False
136                    time.sleep(1)
137                    continue
138            self.run_action(action)
139
140        self.exe.cur.connection.close()
141        if self.exe.ws:
142            self.exe.ws.close()
143
144    def run_action(self, action: Action) -> None:
145        assert self.exe
146        try:
147            if action.run(self.exe):
148                self.num_queries[type(action)] += 1
149                self.num_successes[type(action)] += 1
150            else:
151                self.num_skips[type(action)] += 1
152        except QueryError as e:
153            self.num_queries[type(action)] += 1
154            # TODO(def-): Reduce number of errors for temp tables/views? At
155            # least the errors will be fast, so maybe not worth it
156            # if "temp" in e.msg:
157            #     print(e.query)
158            #     print(e.msg)
159            for error_to_ignore in action.errors_to_ignore(self.exe):
160                if error_to_ignore in e.msg:
161                    self.ignored_errors[error_to_ignore][type(action)] += 1
162                    if (
163                        "Please disconnect and re-connect" in e.msg
164                        or "server closed the connection unexpectedly" in e.msg
165                        or "Can't create a connection to host" in e.msg
166                        or "Connection refused" in e.msg
167                        or "the connection is lost" in e.msg
168                        or "connection in transaction status INERROR" in e.msg
169                    ):
170                        self.exe.reconnect_next = True
171                    else:
172                        self.exe.rollback_next = True
173                    break
174            else:
175                thread_name = threading.current_thread().getName()
176                self.occurred_exception = e
177                print(f"+++ [{thread_name}] Query failed: {e.query} {e.msg}")
178                raise
179        except Exception as e:
180            self.occurred_exception = e
181            raise e

class Worker: View Source

 31class Worker:
 32    rng: random.Random
 33    action_list: ActionList | None
 34    actions: list[Action]
 35    weights: list[float]
 36    end_time: float
 37    num_queries: Counter[type[Action]]
 38    num_successes: Counter[type[Action]]
 39    num_skips: Counter[type[Action]]
 40    autocommit: bool
 41    system: bool
 42    exe: Executor | None
 43    ignored_errors: defaultdict[str, Counter[type[Action]]]
 44    composition: Composition | None
 45    occurred_exception: Exception | None
 46
 47    def __init__(
 48        self,
 49        rng: random.Random,
 50        actions: list[Action],
 51        weights: list[float],
 52        end_time: float,
 53        autocommit: bool,
 54        system: bool,
 55        composition: Composition | None,
 56        action_list: ActionList | None = None,
 57    ):
 58        self.rng = rng
 59        self.action_list = action_list
 60        self.actions = actions
 61        self.weights = weights
 62        self.end_time = end_time
 63        self.num_queries = Counter()
 64        # Unlike num_queries, these are never cleared: they feed the
 65        # end-of-run action coverage check.
 66        self.num_successes = Counter()
 67        self.num_skips = Counter()
 68        self.autocommit = autocommit
 69        self.system = system
 70        self.ignored_errors = defaultdict(Counter)
 71        self.composition = composition
 72        self.occurred_exception = None
 73        self.exe = None
 74
 75    def run(
 76        self, host: str, pg_port: int, http_port: int, user: str, database: Database
 77    ) -> None:
 78        # In scenarios with kills and deploys materialized can go down at any
 79        # point during the setup, keep retrying.
 80        for i in range(300):
 81            try:
 82                self.conn = psycopg.connect(
 83                    host=host, port=pg_port, user=user, dbname="materialize"
 84                )
 85                self.conn.autocommit = self.autocommit
 86                cur = self.conn.cursor()
 87                ws = websocket.WebSocket()
 88                ws_conn_id, ws_secret_key = ws_connect(ws, host, http_port, user)
 89                self.exe = Executor(self.rng, cur, ws, database, user=user)
 90                self.exe.set_isolation("SERIALIZABLE")
 91                cur.execute("SET auto_route_catalog_queries TO false")
 92                if self.exe.use_ws:
 93                    self.exe.pg_pid = ws_conn_id
 94                else:
 95                    cur.execute("SELECT pg_backend_pid()")
 96                    self.exe.pg_pid = cur.fetchall()[0][0]
 97            except Exception:
 98                if time.time() > self.end_time:
 99                    return
100                if i == 299:
101                    raise
102                time.sleep(1)
103            else:
104                break
105        assert self.exe
106
107        while time.time() < self.end_time:
108            action = self.rng.choices(self.actions, self.weights)[0]
109            if not action.applicable(self.exe):
110                continue
111            if self.exe.rollback_next:
112                try:
113                    self.exe.rollback()
114                except QueryError:
115                    # ROLLBACK can itself fail, e.g. cancelled by
116                    # `pg_cancel_backend` or on a broken WS session. Force a
117                    # reconnect rather than leaving a session with an open
118                    # aborted transaction behind.
119                    self.exe.reconnect_next = True
120                    self.exe.rollback_next = False
121                    continue
122                self.exe.rollback_next = False
123            if self.exe.reconnect_next:
124                self.exe.reconnect_next = False
125                # Run as its own action so failures are attributed to
126                # ReconnectAction, not to `action`, which hasn't run yet.
127                self.run_action(
128                    ReconnectAction(self.rng, self.composition, random_role=False)
129                )
130                if self.exe.reconnect_next or self.exe.rollback_next:
131                    # Reconnecting failed with an ignored error. Always retry
132                    # the reconnect, never fall through to the action: the
133                    # old session may hold an aborted transaction that fails
134                    # all statements.
135                    self.exe.reconnect_next = True
136                    self.exe.rollback_next = False
137                    time.sleep(1)
138                    continue
139            self.run_action(action)
140
141        self.exe.cur.connection.close()
142        if self.exe.ws:
143            self.exe.ws.close()
144
145    def run_action(self, action: Action) -> None:
146        assert self.exe
147        try:
148            if action.run(self.exe):
149                self.num_queries[type(action)] += 1
150                self.num_successes[type(action)] += 1
151            else:
152                self.num_skips[type(action)] += 1
153        except QueryError as e:
154            self.num_queries[type(action)] += 1
155            # TODO(def-): Reduce number of errors for temp tables/views? At
156            # least the errors will be fast, so maybe not worth it
157            # if "temp" in e.msg:
158            #     print(e.query)
159            #     print(e.msg)
160            for error_to_ignore in action.errors_to_ignore(self.exe):
161                if error_to_ignore in e.msg:
162                    self.ignored_errors[error_to_ignore][type(action)] += 1
163                    if (
164                        "Please disconnect and re-connect" in e.msg
165                        or "server closed the connection unexpectedly" in e.msg
166                        or "Can't create a connection to host" in e.msg
167                        or "Connection refused" in e.msg
168                        or "the connection is lost" in e.msg
169                        or "connection in transaction status INERROR" in e.msg
170                    ):
171                        self.exe.reconnect_next = True
172                    else:
173                        self.exe.rollback_next = True
174                    break
175            else:
176                thread_name = threading.current_thread().getName()
177                self.occurred_exception = e
178                print(f"+++ [{thread_name}] Query failed: {e.query} {e.msg}")
179                raise
180        except Exception as e:
181            self.occurred_exception = e
182            raise e

Worker( rng: random.Random, actions: list[materialize.parallel_workload.action.Action], weights: list[float], end_time: float, autocommit: bool, system: bool, composition: materialize.mzcompose.composition.Composition | None, action_list: materialize.parallel_workload.action.ActionList | None = None) View Source

47    def __init__(
48        self,
49        rng: random.Random,
50        actions: list[Action],
51        weights: list[float],
52        end_time: float,
53        autocommit: bool,
54        system: bool,
55        composition: Composition | None,
56        action_list: ActionList | None = None,
57    ):
58        self.rng = rng
59        self.action_list = action_list
60        self.actions = actions
61        self.weights = weights
62        self.end_time = end_time
63        self.num_queries = Counter()
64        # Unlike num_queries, these are never cleared: they feed the
65        # end-of-run action coverage check.
66        self.num_successes = Counter()
67        self.num_skips = Counter()
68        self.autocommit = autocommit
69        self.system = system
70        self.ignored_errors = defaultdict(Counter)
71        self.composition = composition
72        self.occurred_exception = None
73        self.exe = None

rng: random.Random

action_list: materialize.parallel_workload.action.ActionList | None

actions: list[materialize.parallel_workload.action.Action]

weights: list[float]

end_time: float

num_queries: collections.Counter[type[materialize.parallel_workload.action.Action]]

num_successes: collections.Counter[type[materialize.parallel_workload.action.Action]]

num_skips: collections.Counter[type[materialize.parallel_workload.action.Action]]

autocommit: bool

system: bool

exe: materialize.parallel_workload.executor.Executor | None

ignored_errors: collections.defaultdict[str, collections.Counter[type[materialize.parallel_workload.action.Action]]]

composition: materialize.mzcompose.composition.Composition | None

occurred_exception: Exception | None

def run( self, host: str, pg_port: int, http_port: int, user: str, database: materialize.parallel_workload.database.Database) -> None: View Source

 75    def run(
 76        self, host: str, pg_port: int, http_port: int, user: str, database: Database
 77    ) -> None:
 78        # In scenarios with kills and deploys materialized can go down at any
 79        # point during the setup, keep retrying.
 80        for i in range(300):
 81            try:
 82                self.conn = psycopg.connect(
 83                    host=host, port=pg_port, user=user, dbname="materialize"
 84                )
 85                self.conn.autocommit = self.autocommit
 86                cur = self.conn.cursor()
 87                ws = websocket.WebSocket()
 88                ws_conn_id, ws_secret_key = ws_connect(ws, host, http_port, user)
 89                self.exe = Executor(self.rng, cur, ws, database, user=user)
 90                self.exe.set_isolation("SERIALIZABLE")
 91                cur.execute("SET auto_route_catalog_queries TO false")
 92                if self.exe.use_ws:
 93                    self.exe.pg_pid = ws_conn_id
 94                else:
 95                    cur.execute("SELECT pg_backend_pid()")
 96                    self.exe.pg_pid = cur.fetchall()[0][0]
 97            except Exception:
 98                if time.time() > self.end_time:
 99                    return
100                if i == 299:
101                    raise
102                time.sleep(1)
103            else:
104                break
105        assert self.exe
106
107        while time.time() < self.end_time:
108            action = self.rng.choices(self.actions, self.weights)[0]
109            if not action.applicable(self.exe):
110                continue
111            if self.exe.rollback_next:
112                try:
113                    self.exe.rollback()
114                except QueryError:
115                    # ROLLBACK can itself fail, e.g. cancelled by
116                    # `pg_cancel_backend` or on a broken WS session. Force a
117                    # reconnect rather than leaving a session with an open
118                    # aborted transaction behind.
119                    self.exe.reconnect_next = True
120                    self.exe.rollback_next = False
121                    continue
122                self.exe.rollback_next = False
123            if self.exe.reconnect_next:
124                self.exe.reconnect_next = False
125                # Run as its own action so failures are attributed to
126                # ReconnectAction, not to `action`, which hasn't run yet.
127                self.run_action(
128                    ReconnectAction(self.rng, self.composition, random_role=False)
129                )
130                if self.exe.reconnect_next or self.exe.rollback_next:
131                    # Reconnecting failed with an ignored error. Always retry
132                    # the reconnect, never fall through to the action: the
133                    # old session may hold an aborted transaction that fails
134                    # all statements.
135                    self.exe.reconnect_next = True
136                    self.exe.rollback_next = False
137                    time.sleep(1)
138                    continue
139            self.run_action(action)
140
141        self.exe.cur.connection.close()
142        if self.exe.ws:
143            self.exe.ws.close()

def run_action(self, action: materialize.parallel_workload.action.Action) -> None: View Source

145    def run_action(self, action: Action) -> None:
146        assert self.exe
147        try:
148            if action.run(self.exe):
149                self.num_queries[type(action)] += 1
150                self.num_successes[type(action)] += 1
151            else:
152                self.num_skips[type(action)] += 1
153        except QueryError as e:
154            self.num_queries[type(action)] += 1
155            # TODO(def-): Reduce number of errors for temp tables/views? At
156            # least the errors will be fast, so maybe not worth it
157            # if "temp" in e.msg:
158            #     print(e.query)
159            #     print(e.msg)
160            for error_to_ignore in action.errors_to_ignore(self.exe):
161                if error_to_ignore in e.msg:
162                    self.ignored_errors[error_to_ignore][type(action)] += 1
163                    if (
164                        "Please disconnect and re-connect" in e.msg
165                        or "server closed the connection unexpectedly" in e.msg
166                        or "Can't create a connection to host" in e.msg
167                        or "Connection refused" in e.msg
168                        or "the connection is lost" in e.msg
169                        or "connection in transaction status INERROR" in e.msg
170                    ):
171                        self.exe.reconnect_next = True
172                    else:
173                        self.exe.rollback_next = True
174                    break
175            else:
176                thread_name = threading.current_thread().getName()
177                self.occurred_exception = e
178                print(f"+++ [{thread_name}] Query failed: {e.query} {e.msg}")
179                raise
180        except Exception as e:
181            self.occurred_exception = e
182            raise e