# Copyright 2018 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import itertools

import math

import queue

import sys

import threading

import tokenize

def orchestrate(*tests, **kwargs):

"""

Orchestrate a deterministic concurrency test.

Runs test functions in separate threads, with each thread taking turns running up

until predefined syncpoints in a deterministic order. All possible orderings are

tested.

Most of the time, we try to use logic, best practices, and static analysis to insure

correct operation of concurrent code. Sometimes our powers of reasoning fail us and,

either through non-determistic stress testing or running code in production, a

concurrent bug is discovered. When this occurs, we'd like to have a regression test

to insure we've understood the problem and implemented a correct solution.

`orchestrate` provides a means of deterministically testing concurrent code so we

can write robust regression tests for complex concurrent scenarios.

`orchestrate` runs each passed in test function in its own thread. Threads then

"take turns" running. Turns are defined by setting syncpoints in the code under

test, using comments containing "pragma: SYNCPOINT". `orchestrate` will scan the

code under test and add syncpoints where it finds these comments.

For example, let's say you have the following code in production::

def hither_and_yon(destination):

hither(destination)

yon(destination)

You've found there's a concurrency bug when two threads execute this code with the

same argument, and you think that by adding a syncpoint between the calls to

`hither` and `yon` you can reproduce the problem in a regression test. First add a

comment with "pragma: SYNCPOINT" to the code under test::

def hither_and_yon(destination):

hither(destination) # pragma: SYNCPOINT

yon(destination)

When testing with orchestrate, there will now be a syncpoint, or a pause, after the

call to `hither` and before the call to `yon`. Now you can write a test to exercise

`hither_and_yon` running in parallel::

from unittest import mock

from tests.unit import orchestrate

from myorg.myproj.sales import travel

def test_concurrent_hither_and_yon():

def test_hither_and_yon():

assert something

travel.hither_and_yon("Raleigh")

assert something_else

counts = orchestrate.orchestrate(test_hither_and_yon, test_hither_and_yon)

assert counts == (2, 2)

What `orchestrate` will do now is take each of the two test functions passed in

(actually the same function, twice, in this case), run them serially, and count the

number of turns it takes to run each test to completion. In this example, it will

take two turns for each test: one turn to start the thread and execute up until the

syncpoint, and then another turn to execute from the syncpoint to the end of the

test. The number of turns will always be one greater than the number of syncpoints

encountered when executing the test.

Once the counts have been taken, `orchestrate` will construct a test sequence that

represents all of the turns taken by the passed in tests, with each value in the

sequence representing the index of the test whose turn it is in the sequence. In

this example, then, it would produce::

[0, 0, 1, 1]

This represents the first test taking both of its turns, followed by the second test

taking both of its turns. At this point this scenario has already been tested,

because this is what was run to produce the counts and the initial test sequence.

Now `orchestrate` will run all of the remaining scenarios by finding all the

permutations of the test sequence and executing those, in turn::

[0, 1, 0, 1]

[0, 1, 1, 0]

[1, 0, 0, 1]

[1, 0, 1, 0]

[1, 1, 0, 0]

You'll notice in our example that since both test functions are actually the same

function, that although it tested 6 scenarios there are effectively only really 3

unique scenarios. For the time being, though, `orchestrate` doesn't attempt to

detect this condition or optimize for it.

There are some performance considerations that should be taken into account when

writing tests. The number of unique test sequences grows quite quickly with the

number of turns taken by the functions under test. Our simple example with two

threads each taking two turns, only yielded 6 scenarios, but two threads each taking

6 turns, for example, yields 924 scenarios. Add another six step thread, for a total

of three threads, and now you have over 17 thousand scenarios. In general, use the

least number of steps/threads you can get away with and still expose the behavior

you want to correct.

For the same reason as above, it is recommended that if you have many concurrent

tests, that you name your syncpoints so that you're not accidentally using

syncpoints intended for other tests, as this will add steps to your tests. While

it's not problematic from a testing standpoint to have extra steps in your tests, it

can use computing resources unnecessarily. A name can be added to any syncpoint

after the `SYNCPOINT` keyword in the pragma definition::

def hither_and_yon(destination):

hither(destination) # pragma: SYNCPOINT hither and yon

yon(destination)

In your test, then, pass that name to `orchestrate` to cause it to use only

syncpoints with that name::

orchestrate.orchestrate(

test_hither_and_yon, test_hither_and_yon, name="hither and yon"

)

As soon as any error or failure is detected, no more scenarios are run

and that error is propagated to the main thread.

One limitation of `orchestrate` is that it cannot really be used with `coverage`,

since both tools use `sys.set_trace`. Any code that needs verifiable test coverage

should have additional tests that do not use `orchestrate`, since code that is run

under orchestrate will not show up in a coverage report generated by `coverage`.

Args:

tests (Tuple[Callable]): Test functions to be run. These functions will not be

called with any arguments, so they must not have any required arguments.

name (Optional[str]): Only use syncpoints with the given name. If omitted, only

unnamed syncpoints will be used.

Returns:

Tuple[int]: A tuple of the count of the number turns for test passed in. Can be

used a sanity check in tests to make sure you understand what's actually

happening during a test.

"""

name = kwargs.pop("name", None)

if kwargs:

raise TypeError(

"Unexpected keyword arguments: {}".format(", ".join(kwargs.keys()))

)

# Produce an initial test sequence. The fundamental question we're always trying to

# answer is "whose turn is it?" First we'll find out how many "turns" each test

# needs to complete when run serially and use that to construct a sequence of

# indexes. When a test's index appears in the sequence, it is that test's turn to

# run. We'll start by constructing a sequence that would run each test through to

# completion serially, one after the other.

test_sequence = []

counts = []

for index, test in enumerate(tests):

thread = _TestThread(test, name)

for count in itertools.count(1): # pragma: NO BRANCH

# Pragma is required because loop never finishes naturally.

thread.go()

if thread.finished:

break

counts.append(count)

test_sequence += [index] * count

# Now we can take that initial sequence and generate all of its permutations,

# running each one to try to uncover concurrency bugs

sequences = iter(_permutations(test_sequence))

# We already tested the first sequence getting our counts, so we can discard it

next(sequences)

# Test each sequence

for test_sequence in sequences:

threads = [_TestThread(test, name) for test in tests]

try:

for index in test_sequence:

threads[index].go()

# Its possible for number of turns to vary from one test run to the other,

# especially if there is some undiscovered concurrency bug. Go ahead and

# finish running each test to completion, if not already complete.

for thread in threads:

while not thread.finished:

thread.go()

except Exception:

# If an exception occurs, we still need to let any threads that are still

# going finish up. Additional exceptions are silently ignored.

for thread in threads:

thread.finish()

raise

return tuple(counts)

_local = threading.local()

class _Conductor:

"""Coordinate communication between main thread and a test thread.

Two way communicaton is maintained between the main thread and a test thread using

two synchronized queues (`queue.Queue`) each with a size of one.

"""

def __init__(self):

self._notify = queue.Queue(1)

self._go = queue.Queue(1)

def notify(self):

"""Called from test thread to let us know it's finished or is ready for its next

turn."""

self._notify.put(None)

def standby(self):

"""Called from test thread in order to block until told to go."""

self._go.get()

def wait(self):

"""Called from main thread to wait for test thread to either get to the

next syncpoint or finish."""

self._notify.get()

def go(self):

"""Called from main thread to tell test thread to go."""

self._go.put(None)

_SYNCPOINTS = {}

"""Dict[str, Dict[str, Set[int]]]: Dict mapping source fileneme to a dict mapping

syncpoint name to set of line numbers where syncpoints with that name occur in the

source file.

"""

def _get_syncpoints(filename):

"""Find syncpoints in a source file.

Does a simple tokenization of the source file, looking for comments with "pragma:

SYNCPOINT", and populates _SYNCPOINTS using the syncpoint name and line number in

the source file.

"""

_SYNCPOINTS[filename] = syncpoints = {}

# Use tokenize to find pragma comments

with open(filename, "r") as pyfile:

tokens = tokenize.generate_tokens(pyfile.readline)

for type, value, start, end, line in tokens:

if type == tokenize.COMMENT and "pragma: SYNCPOINT" in value:

name = value.split("SYNCPOINT", 1)[1].strip()

if not name:

name = None

if name not in syncpoints:

syncpoints[name] = set()

lineno, column = start

syncpoints[name].add(lineno)

class _TestThread:

"""A thread for a test function."""

thread = None

finished = False

error = None

at_syncpoint = False

def __init__(self, test, name):

self.test = test

self.name = name

self.conductor = _Conductor()

def _run(self):

sys.settrace(self._trace)

_local.conductor = self.conductor

try:

self.test()

except Exception as error:

self.error = error

finally:

self.finished = True

self.conductor.notify()

def _sync(self):

# Tell main thread we're finished, for now

self.conductor.notify()

# Wait for the main thread to tell us to go again

self.conductor.standby()

def _trace(self, frame, event, arg):

"""Argument to `sys.settrace`.

Handles frames during test run, syncing at syncpoints, when found.

Returns:

`None` if no more tracing is required for the function call, `self._trace`

if tracing should continue.

"""

if self.at_syncpoint:

# We hit a syncpoint on the previous call, so now we sync.

self._sync()

self.at_syncpoint = False

filename = frame.f_globals.get("__file__")

if not filename:

# Can't trace code without a source file

return

if filename.endswith(".pyc"):

filename = filename[:-1]

if filename not in _SYNCPOINTS:

_get_syncpoints(filename)

syncpoints = _SYNCPOINTS[filename].get(self.name)

if not syncpoints:

# This file doesn't contain syncpoints, don't continue to trace

return

# We've hit a syncpoint. Execute whatever line the syncpoint is on and then

# sync next time this gets called.

if frame.f_lineno in syncpoints:

self.at_syncpoint = True

return self._trace

def go(self):

if self.finished:

return

if self.thread is None:

self.thread = threading.Thread(target=self._run)

self.thread.start()

else:

self.conductor.go()

self.conductor.wait()

if self.error:

raise self.error

def finish(self):

while not self.finished:

try:

self.go()

except Exception:

pass

class _permutations:

"""Generates a sequence of all permutations of `sequence`.

Permutations are returned in lexicographic order using the "Generation in

lexicographic order" algorithm described in `the Wikipedia article on "Permutation"

<https://en.wikipedia.org/wiki/Permutation>`_.

This implementation differs significantly from `itertools.permutations` in that the

value of individual elements is taken into account, thus eliminating redundant

orderings that would be produced by `itertools.permutations`.

Args:

sequence (Sequence[Any]): Sequence must be finite and orderable.

Returns:

Sequence[Sequence[Any]]: Set of all permutations of `sequence`.

"""

def __init__(self, sequence):

self._start = tuple(sorted(sequence))

def __len__(self):

"""Compute the number of permutations.

Let the number of elements in a sequence N and the number of repetitions for

individual members of the sequence be n1, n2, ... nx. The number of unique

permutations is: N! / n1! / n2! / ... / nx!.

For example, let `sequence` be [1, 2, 3, 1, 2, 3, 1, 2, 3]. The number of unique

permutations is: 9! / 3! / 3! / 3! = 1680.

See: "Permutations of multisets" in `the Wikipedia article on "Permutation"

<https://en.wikipedia.org/wiki/Permutation>`_.

"""

repeats = [len(list(group)) for value, group in itertools.groupby(self._start)]

length = math.factorial(len(self._start))

for repeat in repeats:

length /= math.factorial(repeat)

return int(length)

def __iter__(self):

"""Iterate over permutations.

See: "Generation in lexicographic order" algorithm described in `the Wikipedia

article on "Permutation" <https://en.wikipedia.org/wiki/Permutation>`_.

"""

current = list(self._start)

size = len(current)

while True:

yield tuple(current)

# 1. Find the largest index i such that a[i] < a[i + 1].

for i in range(size - 2, -1, -1):

if current[i] < current[i + 1]:

break

else:

# If no such index exists, the permutation is the last permutation.

return

# 2. Find the largest index j greater than i such that a[i] < a[j].

for j in range(size - 1, i, -1):

if current[i] < current[j]:

break

else: # pragma: NO COVER

raise RuntimeError("Broken algorithm")

# 3. Swap the value of a[i] with that of a[j].

temp = current[i]

current[i] = current[j]

current[j] = temp

# 4. Reverse the sequence from a[i + 1] up to and including the final

# element a[n].

current = current[: i + 1] + list(reversed(current[i + 1 :]))