DEV Community: Andrew Werner

Composing generic data structures in go

Andrew Werner — Tue, 30 Nov 2021 14:57:38 +0000

Update 12/1/2021: @rogpeppe gave great feedback on the code and post, it has been updated accordingly.

tl;dr

github.com/ajwerner/btree provides go1.18 generic implementations of CoW btree data structures including regular a regular map/set, order-statistic trees, and interval trees all backed by a shared, generic augmented btree implementation. Given go 1.18 hasn't been released yet and it hasn't been heavily tested, don't use it yet. Benchmarks and a deeper discussion on the patterns and experience to follow.

Potentially motivating context (feel free to skip)

Recently a colleague, Nathan, reflecting on CockroachDB, remarked (paraphrased from memory) that the key data structure is the interval btree. The story of Nathan’s addition of the first interval btree to cockroach and the power of copy-on-write data structures
is worthy of its own blog post for another day. It’s Nathan’s hand-specialization of that data structure that provided the basis (and tests) for the generalization I’ll be presenting here. The reason for this specialization was as much for the performance wins of avoiding excessive allocations, pointer chasing, and cost of type assertions when using interface boxing.

I won't quarrel with the idea that interval btrees are critical, but there's other problems out there which benefit from different types of ordered collection data structures, like regular ol' sorted sets and maps. I'm not pleased with what's currently on offer for go in that space. Anyway, an interval tree is just one of a family of augmented tree search data structures. Another useful augmented tree data structure is the order-statistic tree. Wouldn't it be great if we could have a solid implementation of all the various search trees?

As we all know, generics are coming in go1.18. So what better time for an experience report building some real things with it? This post will explore github.com/ajwerner/btree, a library leveraging go1.18's parametric polymorphism to build an abstract augmented btree which offers rich core functionality and easy extensibility.

The post will go through a high-level overview of the library. Some follow-up posts will cover benchmarks, iterator patterns, and some gripes.

The library

The root of the module is a vanilla sorted set and map library. Here's an example:

m := btree.MakeMap[string, int](strings.Compare)
m.Upsert("foo", 1)
m.Upsert("bar", 2)
fmt.Println(m.Get("foo"))
fmt.Println(m.Get("baz"))
it := m.Iterator()
for it.First(); it.Valid(); it.Next() {
    fmt.Println(it.Cur(), it.Value())
}
// Output:
// 1 true
// 0 false
// foo 1
// bar 2

Let's have a look at this Map:

// Map is a ordered map from K to V.
type Map[K, V any] struct {
    abstract.Map[K, V, struct{}]
}

Here we see that we embed something called abstract.Map. This abstract.Map is a generalization of a sorted map that also allows for extra state as specified by its final type parameter. In this case, for a regular btree, it's set to struct{} is the auxiliary configuration state for the tree.

This augmentation parameter allows for implementers to provide some extra state which will get plumbed down into both the augmentation and the iterator. This is useful for, say, comparison functions or other functions to manipulate the key data. The last two type parameters refer to the augmentation which lives in each key node. Let's have a look at the implementation:

Package `orderstat`

Now, how do we use this augmentation? Armed with this generalization, we can build fancy things like an order-statistic tree with very little effort. This ordered.Compare thing is just a little generic helper to make a comparison function for any type.

s := orderstat.MakeSet(ordered.Compare[int])
for _, i := range rand.Perm(100) {
    s.Upsert(i)
}
fmt.Println(s.Len())
it := m.Iterator()
it.SeekNth(90)
it.Println(s.Cur())
// Output:
// 100
// 90

So, what does this augmentation look like? This:

// Set is an ordered set with items of type T which additionally offers the
// methods of an order-statistic tree on its iterator.
type Set[T any] Map[T, struct{}]

// Map is a ordered map from K to V which additionally offers the methods
// of a order-statistic tree on its iterator.
type Map[K, V any] struct {
    abstract.Map[K, V, aug]
}

type aug struct {
    // children is the number of items rooted at the current subtree.
    children int
}

Now, in order to make this useful, we implement the Updater[K, A] interface which can be used to update the aug method on the aug and we add some special behavior to the Iterator. In order to empower the orderstat library to do something useful with its iterator, we need to expose low-level iterator operations to it. This is done through a constructor in the internal/abstract library which can be used to trade a regular Iterator for a LowLevelIterator. You can see this in action here.

The Updater interface is used to update the augmented data when the node is modified. The UpdateMeta has extra information which can be used to optimize the update.

// Updater is used to update the augmentation of the node when the subtree
// changes.
type Updater[K, V, A any] interface {

    // Update should update the augmentation of the passed node, optionally
    // using the data in the UpdataInfo to optimize the update. If the
    // augmentation changed, and thus, changes should occur in the ancestors
    // of the subtree rooted at this node, return true.
    Update(*Node[K, V, A], UpdateInfo[K, A]) (changed bool)
}

`interval`

Another cool augmented search tree is the interval tree which lets you search for overlaps.

type pair[T constraints.Ordered] [2]T
func (p pair[T]) first() T
func (p pair[T]) second() T
func (p pair[T]) compare(o pair[T]) int {
    if c := ordered.Compare(p[0], o[0]); c != 0 {
        return c
    }
    return ordered.Compare(p[1], o[1])
}

m := interval.MakeSet(
    ordered.Compare[int],
    pair[int].compare,
    pair[int].first,
    pair[int].second,
    nil, // used to determine whether an interval is a point
)
for _, p := range []pair{
    {1, 2}, {2, 3}, {1, 5}, {0, 6}, {2, 7},
} {
    m.Upsert(p)
}
it := m.Iterator()
for it.FirstOverlap(pair{4, 5}); it.Valid(); it.NextOverlap() {
    fmt.Println(it.Cur())
}
// Output:
// [0 6]
// [1 5]
// [2 7]

Recap

We've now got a set of search tree structure which are generic, zero-allocation and share a common core. Each of them offer an O(1) Clone() operation, which is very handy. The business logic related to the augmented search doesn't get mixed up with the business logic of the tree itself. The performance is on par with the hand-specialized interval variant (benchmarks coming later). Expect a follow-up post at some point about rough edges and patterns in the implementation. All in all, I'm excited about the prospect of using this in 2022 in real code.

Special Thanks

A big thank you to Nathan VanBenschoten (@natevanben) who wrote the wonderful code that inspired everything here.

Also, shout out to Roger Peppe (@rogpeppe) who helped me grapple
with the type sets change in the go2go days and helped make this library actually good after I published the first draft of it.

Quick and Easy, Exactly-Once, Distributed Work Queues Using Serializable Transactions

Andrew Werner — Fri, 01 May 2020 18:50:47 +0000

Coordinating work across multiple processes almost always pops up in one form or another when building applications. There are well-known recipes to perform this coordination on top of both distributed coordination services like etcd or Apache Zookeeper/Curator, and single-machine SQL databases like Postgres. Given the strong consistency of CockroachDB, similar patterns always seemed possible but weren't well documented. This short post tries to provide a recipe to achieve correct, fault-tolerant, distributed coordination on top of CockroachDB.

I want to clarify that providing true execute-exactly-once semantics or mutual exclusion for work outside of the coordinating storage system is impossible to provide in a fault-tolerant, highly available way (see Section 3.2 in the Jepson etcd evaluation). In an arbitrarily asynchronous network where nodes may pause for unlimited amounts of time, we can’t know whether an operation is failed or slow.

Fortunately, we can get pretty close by making some reasonable assumptions about failure detection. Secondly, we can provide a different guarantee: exactly once acknowledgement. With this guarantee, only one task will ever be able to successfully acknowledge completion of some work item. If the completion of the task means writing a result to the database, this means we can achieve exactly-once semantics 🎉.

The key concept that enables coordination over work items is the session. The session encapsulates coordinated failure detection. The system will refer to sessions when attempting to perform work. Work items can only be successfully consumed by workers which have a live session throughout the entire processing of a work item.

Imagine we have work items that we want to acknowledge having worked on exactly once. To do this we’ll need two tables, a sessions table and an items table. The trick is that we’re going to use the sessions table to track the liveness of our workers. Each worker will heartbeat its session to indicate its liveness. Workers will lock work items prior to processing by associating the item with their session.

We’ll have workers do 3 things continuously:

1) Heartbeat the sessions table to keep their session alive.
2) Poll the sessions table to see if there are any expired sessions which need to be deleted.
3) Poll the items table to claim work items.

For each work item which is claimed, we'll then process the work item, then run a transaction which records its outcome and removes it from the queue. The code will ensure, transactionally, that the session under which the item was claimed is still valid when recording the result.

Sessions and Heartbeating

CREATE TABLE sessions (
    id
        UUID DEFAULT gen_random_uuid() PRIMARY KEY,
    heartbeated_at
        TIMESTAMPTZ NOT NULL
    -- Probably also a good idea to put some metadata about the session
    -- e.g. the hostname and pid which owns the session and maybe the
    -- timestamp at which it was created.
);

When a process starts up, it will need to create a session:

INSERT INTO sessions (heartbeated_at) VALUES (now()) RETURNING id

Then, the session will need to kick off a thread to periodically update the heartbeated_at timestamp with a new value of now. This should happen generally quite a bit more frequently than the expiration (which you’ll also need to pick).

-- Here $1 is the session ID for the process.
UPDATE sessions SET heartbeated_at = now() WHERE id = $1

If the process fails to heartbeat, likely because your session has expired, then the code should exit the worker process as it won’t be able to acknowledge any work that is currently claimed. This is the case where the system thinks your worker session has failed.

Removing Expired Sessions

We want to ensure that live sessions are eventually able to claim work items that had previously been claimed by now dead sessions. We do this by detecting expired sessions and deleting them.

-- Here $1 is the chosen expiration as an INTERVAL.
DELETE FROM sessions WHERE heartbeated_at < now() - $1

Reasonable values for the heartbeat interval and expiration might be 1 and 5 seconds respectively. With these settings, it would take 5 seconds for a failed node's items to be claimed by another node. Note that during graceful shutdown, the process can remove its session cooperatively, before it has expired, meaning work can be claimed immediately.

Claiming and Processing Work

We'll need an items table to store work items. We could use an index to order the items or just grab them arbitrarily.

CREATE TABLE items (
    id
        UUID NOT NULL PRIMARY KEY,
    claim
        UUID REFERENCES sessions (id) ON DELETE SET NULL
    -- other fields related to score, details of the item, etc.
    -- You could use score to prioritize item. 
    -- Maybe time added and then order by that? 
    -- Also you could create an index.
);

In our work claiming loop, while we have capacity to grab outstanding work items, we'll ask the database for a new item by periodically polling. The polling rate should be chosen based on the latency requirements of the application. Another consideration is the total number of worker nodes. If more nodes are polling for items, consider setting the polling rate to be less frequent. Something on the order of 100-200ms seems pretty reasonable for lots of applications with 10's of worker nodes.

We'll use the below query to find a work item and claim it. We could add an ORDER BY clause to prioritize work items based on properties.

-- Here $1 is the process's session ID.
UPDATE items SET claim = $1 WHERE claim IS NULL LIMIT 1 RETURNING id

Here’s the magic to achieve exactly once semantics: in the processing of this item, you will use a transaction that records the outcome of the item and also deletes the item from the queue. It’s critical that when that transaction runs, it ensures that the process’s session still has the claim on the item. That looks like:

-- do some processing of the item.
BEGIN;

--- ... Record the results of processing the item.

-- Here $1 is the item id and $2 is the process's session.
DELETE FROM items WHERE id = $1 AND claim = $2;
-- SQL clients give access to the number of rows affected by a statement.
-- Ensure that the DELETE affected 1 row. If it didn't, then your session
-- has lost its claim and the transaction needs to be aborted with a
-- ROLLBACK. Otherwise, proceed to COMMIT.

COMMIT;

If the above transaction commits, because of the properties of serializable isolation and transactions, you’ll be sure that the item was processed exactly once!