Aadhav Vignesh

Building an open-source version of Antithesis, Part 1: Understanding the deterministic simulation testing ecosystem

2024-06-14T00:00:00Z

Understanding deterministic simulation testing and laying out the foundations for our implementation

Introduction ¶

You might be wondering, what’s Antithesis? This is what they describe themselves as (from their website):

Antithesis is a continuous reliability platform that autonomously searches for problems in your software within a simulated environment. Every problem we find can be perfectly reproduced, allowing for efficient debugging of even the most complex problems.

The people behind Antithesis are the ones who pioneered DST (deterministic simulation testing) at FoundationDB. But what’s deterministic simulation testing?

To some, DST is a superpower. But for others, it just looks like some variant of automated testing. I’ll take a neutral stance here, and I’ll try to explain it without hyping it up.

As a developer, you likely won’t have the superpower to predict what might break your code. Not everything is observable, and it is very likely that your code might encounter a few “unknown-unknowns” in its lifetime.

Most systems are non-deterministic in nature. Every execution of your program is affected by external factors like clocks/timers, schedulers, networks, certain CPU instructions, etc., which introduce non-determinism.

Testing these systems can be a challenge. Traditional testing methods like unit/integration tests are limited in their ability to reproduce bugs. Some teams might have the resources to model and verify their systems through formal verification, but others likely won’t have the time and resources to follow the path of formal verification.

This is where DST (deterministic simulation testing) steps in. The best case study is FoundationDB, which I’ll be covering later in this article.

To understand DST, I encourage you to listen to Antithesis’ founder Will Wilson’s talk at StrangeLoop 2014, because he explains it better than I ever could.

But if you’re lazy, here’s a short description: DST involves running your code on a single-threaded simulator, where most external interfaces that introduce non-determinism are mocked to make it deterministic. The simulator then tries to find bugs through injecting multiple faults, and these bugs can be reproduced deterministically based on the seed used for the run.

Isn’t this fault injection or chaos engineering? Well.. technically yes (to some extent), but where DST shines is that it allows for advancing time just like how time advancement works in DES (discrete event simulation). Discrete-event simulation (DES) models the operations of a system as a sequence of events in time, where each event occurs at a particular point in time, and marks a change of state in the system. One could say DST is an extended version of DES, but I’ll leave the interpretation up to the reader.

With the advantage of time advancement, you can simulate years of system behavior in shorter periods of time, which makes it highly effective for understanding system behavior and uncovering issues that become visible after a certain period of time has elapsed.

Another advantage of utilizing deterministic simulation testing is that since most interfaces like networks are mocked, this reduces the amount of time required to process that operation.¹

But the catch here is that your system needs to be deterministic to utilize deterministic simulation testing properly. FoundationDB, TigerBeetle and others have been able to utilize DST due to their design decisions, as they focused on making their systems deterministic.

Not everyone has the time or interest in rewriting their systems to be deterministic in nature. And some don’t want to deal with the overhead that DST brings, in terms of increased integration effort and design challenges, and I feel that it’s completely fine given that we live in an era where most systems are getting pointless integrations (psst, AI).

But what about those who want to integrate DST in their workloads but don’t have the resources to do so?

Antithesis tries to solve this problem. People who want to reap the benefits of deterministic simulation testing can do it with minimal changes to their current systems with the Antithesis platform. To use Antithesis, you just need to package your apps as containers, add some configuration and a workload to test on, and let Antithesis run the tests². For obvious reasons, it’s a paid platform, but I’m not sure if it’s affordable enough for the general audience.

Why do I want to build one? ¶

But even before all that, you might be thinking about the goal of this project. Why do I even want to work on an open-source version?

There are two primary reasons behind it:

As a learning exercise. Antithesis has a team full of talented engineers who have experience working in this space for almost 10 years. I definitely won’t be able to come close to their implementation, but even if I come close to a working prototype, I’d say that would be a huge achievement for me.
No individual sign-up available.
- Antithesis currently doesn’t have an option to sign up as a single user, the only way right now to sign up is to contact their sales team. (However, if you’re a maintainer of an open-source project, they’ll try to test your code for free. 3)
- So, I decided that building an open-source version would be the best bet, and even though it’s easier said than done, it’s better to actively work on an OSS/free version for competitiveness.

I know that this is an ambitious project, and it’s likely that a working prototype would take years of development. But I also want to try experimenting without worrying about the fact that I’m trying to emulate the efforts of a team of experienced engineers who took ~6 years to build Antithesis.

For now, as a starting point, I’m trying to understand the various systems that utilize deterministic simulation testing. I’ll discuss the ones I know below in some detail.

FoundationDB (popular!) ¶

If you’re reading this article, then you might already know about FoundationDB. But if you don’t, then you might be interested in knowing their approach towards building resilient systems.

FoundationDB is a fault-tolerant key-value database which got acquired by Apple in 2015. It’s widely known for its emphasis on testing which is described in detail here and in their paper. I highly recommend reading both! (those are some of my favorites in my reading list)

To summarize, they first built an actor model on top of C++ called Flow (even before they wrote their database). Through Flow, they were able to abstract various actions of the database into a bunch of actors, which can be scheduled by Flow’s runtime. As a result of modeling their logic into an actor-based programming model, they were able to easily integrate it with their deterministic simulator. There’s a lot more to it, and I’ll try to cover it in the upcoming posts.

Here’s a HN comment (2013) from the founder about how Flow came into existence before FoundationDB:

We knew this was going to be a long project so we invested heavily in tools at the beginning. The first two weeks of FoundationDB were building this new programming language to give us the speed of C++ with high level tools for actor-model concurrency. But, the real magic is how Flow enables us to use our real code to do deterministic simulations of a cluster in a single thread.

As the first step³, the simulator will initialize a random seed like this (if not provided through the CLI):

// https://github.com/apple/foundationdb/blob/f27bc4ac2ba78de23e062ecfb3e8bc9a304e0c6e/fdbserver/fdbserver.actor.cpp#L1070

uint32_t randomSeed = platform::getRandomSeed();

The simulator will then try to run the specified workload inside the simulated cluster where network, disk, time, etc. are mocked to remove non-determinism. Here’s how a workload (or a test file) looks like:

# https://github.com/apple/foundationdb/blob/f27bc4ac2ba78de23e062ecfb3e8bc9a304e0c6e/tests/fast/AtomicBackupCorrectness.toml

[[test]]
testTitle = 'BackupAndRestore'
clearAfterTest = false
simBackupAgents = 'BackupToFile'

    [[test.workload]]
    testName = 'AtomicOps'
    nodeCount = 30000
    transactionsPerSecond = 2500.0
    testDuration = 30.0

    [[test.workload]]
    testName = 'BackupAndRestoreCorrectness'
    backupAfter = 10.0
    restoreAfter = 60.0
    backupRangesCount = -1

    [[test.workload]]
    testName = 'RandomClogging'
    testDuration = 90.0

    [[test.workload]]
    testName = 'Rollback'
    meanDelay = 90.0
    testDuration = 90.0

    [[test.workload]]
    testName = 'Attrition'
    machinesToKill = 10
    machinesToLeave = 3
    reboot = true
    testDuration = 90.0

    [[test.workload]]
    testName = 'Attrition'
    machinesToKill = 10
    machinesToLeave = 3
    reboot = true
    testDuration = 90.0

This is a test composed of smaller workloads like RandomClogging (packet transfer delays/pauses, or simply, messing with the network), Attrition (kill/reboot machines), etc.

Since the simulator is deterministic, you can use the same seed which was generated at the start of the test to reproduce the bug that occurred.

The FoundationDB Simulator (source: paper)

Another strategy used by FoundationDB is to extensively use their BUGGIFY macros⁴, which helps the simulator in finding more interesting edge cases which might lead to bugs. These macros are only enabled when they’re run in the simulator.

A few examples can be seen here, this one tries to tweak server knobs:

// https://github.com/apple/foundationdb/blob/f27bc4ac2ba78de23e062ecfb3e8bc9a304e0c6e/fdbclient/ServerKnobs.cpp#L556

init( ROCKSDB_ENABLE_CHECKPOINT_VALIDATION,                false ); if ( randomize && BUGGIFY ) ROCKSDB_ENABLE_CHECKPOINT_VALIDATION = deterministicRandom()->coinflip();

and here the BUGGIFY macro is used inside a test workload for randomly closing connections:

// https://github.com/apple/foundationdb/blob/main/fdbserver/workloads/HTTPKeyValueStore.actor.cpp#L286-L291

// sometimes randomly close connection anyway
if (BUGGIFY_WITH_PROB(0.1)) {
    ASSERT(self->conn.isValid());
    self->conn->close();
    self->conn.clear();
}

I personally like this approach of introducing failures through code and assisting the simulator, instead of letting the simulator trying to find interesting execution paths on its own.

TigerBeetle ¶

TigerBeetle’s simulator called VOPR (Viewstamped Operation Replicator)⁵, tries to spawn a cluster consisting of multiple servers and clients interacting with each other. All types of I/O is mocked out, and the entire simulation runs as a single process.

Network faults and delays are introduced through the network simulator. Here’s an example of network fault injection:

// https://github.com/tigerbeetle/tigerbeetle/blob/0277b9bf4e29443e12bae4cfed36f8306c721ef0/src/testing/packet_simulator.zig#L366

if (self.options.node_count > 1 and self.should_partition()) {
    self.auto_partition_network();

Storage faults are introduced through the in-memory storage simulator. Here’s how TigerBeetle tries to corrupt sectors:

// https://github.com/tigerbeetle/tigerbeetle/blob/0277b9bf4e29443e12bae4cfed36f8306c721ef0/src/testing/storage.zig#L222

/// Cancel any currently in-progress reads/writes.
/// Corrupt the target sectors of any in-progress writes.
pub fn reset(storage: *Storage) void {
    log.debug("Reset: {} pending reads, {} pending writes, {} pending next_ticks", .{
        storage.reads.len,
        storage.writes.len,
        storage.next_tick_queue.count,
    });
    while (storage.writes.peek()) |_| {
        const write = storage.writes.remove();
        if (!storage.x_in_100(storage.options.crash_fault_probability)) continue;

        // Randomly corrupt one of the faulty sectors the operation targeted.
        // TODO: inject more realistic and varied storage faults as described above.
        const sectors = SectorRange.from_zone(write.zone, write.offset, write.buffer.len);
        storage.fault_sector(write.zone, sectors.random(storage.prng.random()));
    }

Another nifty feature is their state checker which checks all state transitions of the replicas. The checksums are based on AEGIS-128L and are constructed here.

// https://github.com/tigerbeetle/tigerbeetle/blob/0277b9bf4e29443e12bae4cfed36f8306c721ef0/src/simulator.zig#L327C9-L331C49

const commits = simulator.cluster.state_checker.commits.items;
const last_checksum = commits[commits.len - 1].header.checksum;
for (simulator.cluster.aofs, 0..) |*aof, replica_index| {
    if (simulator.core.isSet(replica_index)) {
        try aof.validate(last_checksum);

And since the VOPR is completely deterministic, we can replay a bug using a seed, similar to the one we saw in FoundationDB.

Turmoil (Tokio) ¶

Tokio announced Turmoil last year and it has been on my radar for a while, but I haven’t seen a lot of projects using it⁶.

Turmoil tries to simulate hosts, networks and time. I’m not sure if they simulate disks, but this is in contrast to the previous approaches (FoundationDB, TigerBeetle, etc.) as the idea here is that you can import turmoil as a Rust crate and write your simulation tests.

The implementation is a bit similar to TigerBeetle’s simulator implementation:

// https://github.com/tokio-rs/turmoil/blob/766108f2e48bc54092955fc374fed2e0a15505f6/src/sim.rs#L137
pub fn crash(&mut self, addrs: impl ToIpAddrs) {
    self.run_with_hosts(addrs, |addr, rt| {
        rt.crash();

        tracing::trace!(target: TRACING_TARGET, addr = ?addr, "Crash");
    });
}

// https://github.com/tokio-rs/turmoil/blob/766108f2e48bc54092955fc374fed2e0a15505f6/src/sim.rs#L367
let World {
    rng,
    topology,
    hosts,
    ..
} = world.deref_mut();
topology.deliver_messages(rng, hosts.get_mut(&addr).expect("missing host"));

There’s a RNG (random number generator), network mocks, etc. embedded into the simulator. On each tick, it performs operations like message delivery, etc. Nothing novel here, but that doesn’t discredit the work done to implement it. I still feel that it’s great of the Tokio team to make this open source in order to help other Rust users.

Coyote ¶

Coyote is another library which can be used to test C# code deterministically, but the crazy part here is that it doesn’t require you to change a single line in your code.

It does this by binary rewriting at test time, which injects code to allow Coyote to take control over the task scheduler. Here’s an example of a rewriting pass which I found to be easily understandable:

// https://github.com/microsoft/coyote/blob/20a461738abb16d595def740fc486c9071a9cbab/Source/Test/Rewriting/Passes/Rewriting/CallSiteExtractionRewritingPass.cs#L14-L50

/// Rewriting pass that injects callbacks to the runtime for extracting call-site information.
/// </summary>
internal sealed class CallSiteExtractionRewritingPass : RewritingPass
{
    /// .... removed some boilerplate here

    /// <inheritdoc/>
    protected override void VisitMethodBody(MethodBody body)
    {
        if (this.IsCompilerGeneratedType || this.IsAsyncStateMachineType ||
            this.Method is null || this.Method.IsConstructor ||
            this.Method.IsGetter || this.Method.IsSetter)
        {
            return;
        }

        // Get the first instruction in the body.
        Instruction nextInstruction = body.Instructions.FirstOrDefault();

        // Construct the instructions for notifying the runtime which method is executing.
        string methodName = GetFullyQualifiedMethodName(this.Method);
        Instruction loadStrInstruction = this.Processor.Create(OpCodes.Ldstr, methodName);

        TypeDefinition providerType = this.Module.ImportReference(typeof(Operation)).Resolve();
        MethodReference notificationMethod = providerType.Methods.FirstOrDefault(m => m.Name == nameof(Operation.RegisterCallSite));
        notificationMethod = this.Module.ImportReference(notificationMethod);
        Instruction callInstruction = this.Processor.Create(OpCodes.Call, notificationMethod);

        this.Processor.InsertBefore(nextInstruction, this.Processor.Create(OpCodes.Nop));
        this.Processor.InsertBefore(nextInstruction, loadStrInstruction);
        this.Processor.InsertBefore(nextInstruction, callInstruction);

During each execution, Coyote tries to find different execution paths through different strategies and generates reproducible traces for the ones which reports bugs. These are the different strategies used (but there might be more):

// https://github.com/microsoft/coyote/blob/20a461738abb16d595def740fc486c9071a9cbab/Source/Core/Runtime/Scheduling/OperationScheduler.cs#L11-L18
using BoundedRandomFuzzingStrategy = Microsoft.Coyote.Testing.Fuzzing.BoundedRandomStrategy;
using DelayBoundingInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.DelayBoundingStrategy;
using DFSInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.DFSStrategy;
using PrioritizationFuzzingStrategy = Microsoft.Coyote.Testing.Fuzzing.PrioritizationStrategy;
using PrioritizationInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.PrioritizationStrategy;
using ProbabilisticRandomInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.ProbabilisticRandomStrategy;
using QLearningInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.QLearningStrategy;
using RandomInterleavingStrategy = Microsoft.Coyote.Testing.Interleaving.RandomStrategy;

Coyote is really neat. I love the less-intrusive approach of code injection, but it’s a bummer that it’s only available for C#.

Madsim ¶

Magical Deterministic Simulator (Madsim) is a deterministic simulator for Rust programs, which follows the same concept as others by utilizing a PRNG (Pseudo-Random Number Generator), a task scheduler, and network/disk simulators which removes the sources of non-determinism from the environment.

To provide deterministic time, it tries to override the libc gettimeofday function:

// https://github.com/madsim-rs/madsim/blob/main/madsim/src/sim/time/system_time.rs#L6-L21

unsafe extern "C" fn gettimeofday(tp: *mut libc::timeval, tz: *mut libc::c_void) -> libc::c_int {
    // NOTE: tz should be NULL.
    // Linux: The use of the timezone structure is obsolete; the tz argument should normally be specified as NULL.
    // macOS: timezone is no longer used; this information is kept outside the kernel.
    if tp.is_null() {
        return 0;
    }
    if let Some(time) = super::TimeHandle::try_current() {
        // inside a madsim context, use the simulated time.
        let dur = time
            .now_time()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap();
        tp.write(libc::timeval {
            tv_sec: dur.as_secs() as _,
            tv_usec: dur.subsec_micros() as _,
        });
// more code ....

They also seem to have took some inspiration from FoundationDB by using BUGGIFY macros for triggering fault injection.

But once again, this can only be used Rust-based projects, but I like the patches made to std API functions which helped them to achieve determinism.

Antithesis ¶

And finally, Antithesis. The most exciting one on this list.

The reason why Antithesis is gaining so much popularity is due to the fact that you can test almost any application deterministically through their platform with minimal changes. This is a great improvement over existing solutions, which are bound to a single language/runtime.

It’s made possible by their proprietary deterministic hypervisor⁷, which ensures that any code which is non-deterministic in nature, can be made deterministic by running it inside the hypervisor.

Their hypervisor is built on top of the bhyve hypervisor with lots of changes made to the core to make it deterministic. I recommend reading their blog post to know about the challenges they faced while building it. It’s an impressive piece of work, and it’s definitely worth the years of effort they’ve put in.

The simulation environment consists of the system under test packaged as a set of containers, along with the workload, running on the virtual machine managed by the deterministic hypervisor.

For finding bugs, they have a “software explorer” which tries to actively find new and interesting execution paths (or, more formally, branch exploration) by exploring the state space efficiently through some guided fuzzing, and by injecting faults through a fault injection framework.

So, to summarize, they seem to follow the same pattern of setting up a deterministic environment (PRNG, mocked I/O, network, etc.) and using fault injection and intelligent fuzzing to explore the state space efficiently to find bugs. But the real differentiator is their deterministic hypervisor, which is a hard one to replicate in my opinion.

Facebook/Meta tried to do a similar thing previously through Hermit, which is a project to execute programs inside a deterministic sandbox, made possible by Reverie, which is used for syscall interception. But Hermit seems to be in maintenance mode, so I’m not sure about its future.

Other components like the “software explorer” also seem to be the shiny parts that make Antithesis a lucrative option, as I assume that their state space exploration is a bit unique (or performant) as compared to modern fuzzing techniques, but I could be completely wrong here, and maybe it’s a good mix of modern and traditional search techniques.

Since Antithesis is a paid product [$], I don’t have much information on its internals, so I can’t speak much about it. And it’s likely that they won’t reveal their secret sauce soon, but I appreciate them for dropping some really interesting demos!

Time to design! ¶

So far, I’ve talked about a few approaches to deterministic simulation testing above, but it’s time to represent the design I have in my mind for this project.

Openthesis⁸ borrows design decisions from the previously discussed systems. Here’s the rough, initial architecture diagram for the project (click to enlarge):

A rough architecture diagram of Openthesis

The system consists of several parts:

The “core”: This is the central orchestrator responsible for setting up clusters based on the workload provided. The workloads are a part of the code, or would be readable inside the container. This is inspired by the FoundationDB approach of testing different scenarios through different workload configurations.
The “explorer”: This piece is responsible for efficient state space exploration. It consists of a “guided fuzzer”, a fault injection tool, and a report generator.
- For the “guided fuzzer”, I’m thinking of utilizing coverage-guided fuzzers like libFuzzer (and maybe even AFL++ and syzkaller?).
- The fault injection tool will assist in the exploration process and its intensity would be tunable through configuration files.
- The report generator will receive the interesting bugs found by the guided fuzzer throughout the simulation’s lifetime and store them in an object store or a database for further processing.
Testbed: This is the deterministic hypervisor where all sources of non-determinism are removed. This is based on Antithesis’ approach, as I’d want this project to be able to test almost all kinds of applications. A huge undertaking though.

Here’s what’s not covered in the diagram:

Code-driven fault-injection: It would be nice to support BUGGIFY-style fault-injection through code, but that would require maintaining SDKs for a few languages. It might be doable, but I haven’t thought about it much right now.
Time travel and snapshots: I’ve looked at the rr project previously, and I’m interested to see how time-travel debugging fits in this picture. But to make this work, VM snapshotting should be quick enough too. I believe I’ll get answers to these questions once I explore more about building deterministic hypervisors.

Conclusion ¶

Here’s the project’s repository. No code yet, but I’ll add a few more design documents (or notes) over the next few weeks.

In part 2, I’ll try to explain a few more things like time-travel debugging (rr project ⁹, etc.), guided fuzzing, and most importantly, how do I want to approach the gigantic task of trying to build the deterministic hypervisor. It’s going to be a real challenge, but I’m excited! And I’m open to any type of contributions: discussions related to design decisions, code, etc.

I’ve intentionally kept the details short for some sections (TigerBeetle, etc.) and I didn’t cover others (Resonate, Dropbox’s randomized testing of Nucleus, and PolarSignals’ WASM-based (mostly) DST) as this post was getting a bit verbose. I recommend checking them out for more details.

If you liked this, please email or tweet me with questions, corrections, or ideas! And if you’re working at Antithesis or similar companies, please do let me know your thoughts!

However, it’s completely dependent on the implementation. A mock may even slow things down. ↩︎
Based on the “getting started” section in the documentation. ↩︎
Well, not really, as there are a few more steps before it reaches there. ↩︎
Alex Miller (ex-FoundationDB) has a great blog on it. ↩︎
They put it inside a browser, and that went viral to some extent on social media. It’s a fun way to showcase the VOPR. ↩︎
Based on my exploration on GitHub, but I might be wrong here. ↩︎
They do claim to make the hypervisor open source, but I’m not sure if they’d ever do it. ↩︎
Openthesis = open-source Antithesis. If you have a better name, please suggest one on the repository! ↩︎
rr is much more than a time-travel debugger, it also helps in providing deterministic simulation for arbitary processes. ↩︎

pg_experiments, Part 2: Adding a new data type

2023-11-01T00:00:00Z

Exploring Postgres internals by adding a new data type.

Disclaimer: This is my second post on my experiment with Postgres internals. I'm still pretty new to this, and I have only explored a few things in the source code, so please let me know if I did things incorrectly.

Edit: I’ve pushed my changes here, feel free to check it out!

Introduction ¶

Two days in, I’ve made some tiny baby steps in understanding how this thing works. In my previous experiment, I was tinkering with + operator logic, so I decided it would be a nice time to explore what it takes to write my own data type.

After doing some research on this, I found out that the easiest way to do it was to develop an extension which would have the logic for the data type. You plug it in into Postgres, and it works as intended. But I wanted to explore further; building an extension didn’t feel like a challenge.

I wanted to add a new “in-built” data type by trying to tweak with Postgres’ backend. Why? Because it’s fun to break things and fix it. I took up the challenge and started finding things out….

Initial exploration ¶

As we saw in the previous experiment, logic for ADTs (abstract data types) is present in src/backend/utils/adt, so it made sense to explore this area once again. In order to not get intimidated by complex logic, I tried to search for an ADT which was easy to follow and had minimal logic.

I found bool.c to be a nice candidate for exploration; it was pretty easy to follow, and gave me an idea about how ADTs are structured. One pattern I found out was that all ADTs seem to implement an interface; most of them had function names with common suffixes like in, out, send, recv. On exploring further, it made sense:

in is meant for converting an external representation to your ADT’s internal representation
out is meant for converting your ADT’s internal representation to an external representation
send is meant for converting your ADT’s internal representation to a binary representation
recv is meant for converting an external binary representation to your ADT’s internal representation

ADTs must always have input and output functions. send and recv are optional and are used for providing binary input and output routines.

For this experiment, I wanted to build a custom color type which would store hex values and if possible, it should support operations on top of it. And since I haven’t looked into the internals of how colors are represented, this was the perfect opportunity.

Implementing the color type ¶

After multiple failed attempts and rewrites, I was finally happy with the internal representation for color - an uint32 integer:

/*-------------------------------------------------------------------------
 *
 * color.c
 *	  PostgreSQL type definition for a custom 'color' data type.
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"
#include "fmgr.h"
#include "utils/builtins.h"
#include "libpq/pqformat.h"

typedef uint32 color;

#define COLOR_RED_SHIFT 16
#define COLOR_GREEN_SHIFT 8
#define COLOR_BLUE_SHIFT 0
#define COLOR_CHANNEL_MASK 0xFF

Next thing to do was to implement the input function. To implement it, I took a cstring as input and converted it to an uint32 by using scanf:

/*
 * color_in - Color reader. Accepts a hexadecimal string as the input, and converts it to a color value.
 */
PG_FUNCTION_INFO_V1(color_in);
Datum color_in(PG_FUNCTION_ARGS)
{
    char *str = PG_GETARG_CSTRING(0);
    uint32 c = 0;

    if (sscanf(str, "%06X", &c) != 1) {
        ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
                        errmsg("invalid input syntax for type color: \"%s\"", str)));
    }

    PG_RETURN_UINT32(c);
}

The output function takes an uint32 as an argument, and converts it to a hexadecimal string using snprintf:

/*
 * color_out - Color output function. Converts the internal color value to a hexadecimal string representation.
 */
PG_FUNCTION_INFO_V1(color_out);
Datum color_out(PG_FUNCTION_ARGS)
{
    uint32 c = PG_GETARG_UINT32(0);
    char str[7];

    // Format color as a 6-character hexadecimal string
    snprintf(str, sizeof(str), "%06X", c);
    PG_RETURN_CSTRING(pstrdup(str));
}

Send and receive routines ¶

Implementing send and recv was tricky. I checked bool’s implementation and it seemed to use a few things from pg_format.c, which had an interesting comment:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/libpq/pqformat.c;h=d712c80b5e76b7b488211291bbb63c51203c7cad;hb=HEAD#l2)

/*-------------------------------------------------------------------------
 *
 * pqformat.c
 *		Routines for formatting and parsing frontend/backend messages
 *
 * Outgoing messages are built up in a StringInfo buffer (which is expansible)
 * and then sent in a single call to pq_putmessage.  This module provides data
 * formatting/conversion routines that are needed to produce valid messages.
 * Note in particular the distinction between "raw data" and "text"; raw data
 * is message protocol characters and binary values that are not subject to
 * character set conversion, while text is converted by character encoding
 * rules.
 *
 * Incoming messages are similarly read into a StringInfo buffer, via
 * pq_getmessage, and then parsed and converted from that using the routines
 * in this module.

This is exactly what I needed to understand. So, I needed to store my message in a buffer of type StringInfo. But wait. Wasn’t the buffer of type StringInfoData in bool.c?

Let me check. Oh, okay.

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/include/lib/stringinfo.h;h=598ed093dc8757fd08f089b0a1c3cb15ec0c37d0;hb=HEAD#l22)

typedef struct StringInfoData
{
	char	   *data;
	int			len;
	int			maxlen;
	int			cursor;
} StringInfoData;

typedef StringInfoData *StringInfo;

The interface routines are also neatly specified in the code:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/libpq/pqformat.c;h=d712c80b5e76b7b488211291bbb63c51203c7cad;hb=HEAD#l32)

/*
 * INTERFACE ROUTINES
 * Message assembly and output:
 *		pq_beginmessage - initialize StringInfo buffer
 *		pq_sendbyte		- append a raw byte to a StringInfo buffer
 *		pq_sendint		- append a binary integer to a StringInfo buffer
 *		pq_sendint64	- append a binary 8-byte int to a StringInfo buffer
 *		pq_sendfloat4	- append a float4 to a StringInfo buffer
 *		pq_sendfloat8	- append a float8 to a StringInfo buffer
 *		pq_sendbytes	- append raw data to a StringInfo buffer
 *		pq_sendcountedtext - append a counted text string (with character set conversion)
 *		pq_sendtext		- append a text string (with conversion)
 *		pq_sendstring	- append a null-terminated text string (with conversion)
 *		pq_send_ascii_string - append a null-terminated text string (without conversion)
 *		pq_endmessage	- send the completed message to the frontend
 * Note: it is also possible to append data to the StringInfo buffer using
 * the regular StringInfo routines, but this is discouraged since required
 * character set conversion may not occur.
 *
 * typsend support (construct a bytea value containing external binary data):
 *		pq_begintypsend - initialize StringInfo buffer
 *		pq_endtypsend	- return the completed string as a "bytea*"
 *
 * Special-case message output:
 *		pq_puttextmessage - generate a character set-converted message in one step
 *		pq_putemptymessage - convenience routine for message with empty body
 *
 * Message parsing after input:
 *		pq_getmsgbyte	- get a raw byte from a message buffer
 *		pq_getmsgint	- get a binary integer from a message buffer
 *		pq_getmsgint64	- get a binary 8-byte int from a message buffer
 *		pq_getmsgfloat4 - get a float4 from a message buffer
 *		pq_getmsgfloat8 - get a float8 from a message buffer
 *		pq_getmsgbytes	- get raw data from a message buffer
 *		pq_copymsgbytes - copy raw data from a message buffer
 *		pq_getmsgtext	- get a counted text string (with conversion)
 *		pq_getmsgstring - get a null-terminated text string (with conversion)
 *		pq_getmsgrawstring - get a null-terminated text string - NO conversion
 *		pq_getmsgend	- verify message fully consumed
 */

For my color type, I needed the following:

pq_begintypsend: for initializing the buffer
pq_endtypsend: for “finalizing” the buffer and converting the buffer contents to bytea
pq_sendint: for appending an integer to the buffer
pq_getmsgint: for getting a binary integer from the buffer

Cool. Let’s use these routines for building our send and receive functionality:

/*
 * color_send - converts color to binary format
 */
PG_FUNCTION_INFO_V1(color_send);
Datum color_send(PG_FUNCTION_ARGS)
{
    uint32 c = PG_GETARG_UINT32(0);
    StringInfoData buf;

    pq_begintypsend(&buf);
    pq_sendint(&buf, c, 4);
    PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}

/*
 * color_recv - converts external binary format back to the internal color type
 */
PG_FUNCTION_INFO_V1(color_recv);
Datum color_recv(PG_FUNCTION_ARGS)
{
    StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
    uint32 c;

    c = pq_getmsgint(buf, 4);
    PG_RETURN_UINT32(c);
}

Figuring out the system catalog ¶

Before testing it out, these functions should be registered in the pg_proc catalog (pg_proc.dat). Entries in pg_proc.dat define the functions associated with your custom data type, specify their return types, and declare the argument types they accept. Here’s what I added to pg_proc.dat for my color type:

{ oid => '7901', proname => 'color_in', prorettype => 'color',
  proargtypes => 'cstring', prosrc => 'color_in' },
{ oid => '7902', proname => 'color_out', prorettype => 'cstring',
  proargtypes => 'color', prosrc => 'color_out' },
{ oid => '7903', proname => 'color_send', prorettype => 'bytea',
  proargtypes => 'color', prosrc => 'color_send' },
{ oid => '7904', proname => 'color_recv', prorettype => 'color',
  proargtypes => 'internal', prosrc => 'color_recv' },

I’ll try to explain what these keys are (as per my understanding):

oid: This field specifies the OID (Object Identifier) of the function. OIDs are unique identifiers for objects in the PostgreSQL system catalog. (I had to find OIDs which aren’t referenced anywhere.) (There is a tool (unused_oids) to help simplify this process!)
proname: The name of the function.
prorettype: It specifies the return data type of the function. In this case, the function color_in returns the color data type.
proargtypes: Defines the argument types of the function. 'cstring' indicates that the input function color_in accepts a single argument of type cstring.
prosrc: The function’s C-language name (link symbol).

And now to register the color type, I added this to the type system catalog (pg_type.dat):

# custom type (color)
{ oid => '7012', array_type_oid => '7201',
  descr => 'color data type',
  typname => 'color', typlen => '4', typbyval => 't', typcategory => 'N',
  typinput => 'color_in', typoutput => 'color_out', typreceive => 'color_recv',
  typsend => 'color_send', typalign => 'i', typstorage => 'x' },

Each of these fields is explained nicely in the documentation. I’ve used typstorage as x because I didn’t know a suitable storage strategy, maybe p is the correct way to do it?

Edit: Thanks to Paul Jungwirth for letting me know that p (plain storage) should be the best choice for typstorage, as the color type is essentially a uint32 and is fixed-width in nature.

Solving compilation errors ¶

I was excited to test it out, so I set up everything:

$ ./configure --prefix=$(pwd)/test-install
$ make -j8 && make install
$ ./test-install/bin/initdb $(pwd)/test-data # Create a database
$ echo "port=8080" > $(pwd)/postgres.conf # This is the minimal postgres run config
$ mkdir test-run # Place for postgres to put lock files
$ ./test-install/bin/postgres --config-file=$(pwd)/postgres.conf -D $(pwd)/test-data -k $(pwd)/test-run

But on compiling, I encountered this error:

Use of uninitialized value in sprintf at genbki.pl line 1027.
unresolved OID reference "color_in" in pg_type.dat field typinput line
Use of uninitialized value in sprintf at genbki.pl line 1027.
unresolved OID reference "color_out" in pg_type.dat field typoutput line
Use of uninitialized value in sprintf at genbki.pl line 1027.
unresolved OID reference "color_recv" in pg_type.dat field typreceive line
Use of uninitialized value in sprintf at genbki.pl line 1027.
unresolved OID reference "color_send" in pg_type.dat field typsend line
make[2]: *** [Makefile:172: bki-stamp] Error 1
make[2]: Leaving directory '/home/burntcarrot/postgresql/src/backend/catalog'
make[1]: *** [Makefile:141: submake-catalog-headers] Error 2
make[1]: Leaving directory '/home/burntcarrot/postgresql/src/backend'
make: *** [src/Makefile.global:385: submake-generated-headers] Error 2

The error is originating from genbki.pl, which is a Perl script which generates postgres.bki:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/catalog/genbki.pl;h=380bc23c82e76ec175efd39100f9f76403c2250d;hb=HEAD#l4)

#!/usr/bin/perl
#----------------------------------------------------------------------
#
# genbki.pl
#    Perl script that generates postgres.bki and symbol definition
#    headers from specially formatted header files and data files.
#    postgres.bki is used to initialize the postgres template database.
#
# Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
# Portions Copyright (c) 1994, Regents of the University of California
#
# src/backend/catalog/genbki.pl
#
#----------------------------------------------------------------------

And what is postgres.bki? On searching through the docs, I found this (link):

To create the catalog files and load this initial data into them, a backend running in bootstrap mode reads a BKI (Backend Interface) file containing commands and initial data. The postgres.bki file used in this mode is prepared from the aforementioned header and data files, while building a PostgreSQL distribution, by a Perl script named genbki.pl.

Wow. A TIL moment. But the problem wasn’t solved yet.

The origin of the error was on line 1027, which seemed to have logic for performing OID lookups:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/catalog/genbki.pl;h=380bc23c82e76ec175efd39100f9f76403c2250d;hb=HEAD#l1025)

# ....

# Perform OID lookups on an array of OID names.
# If we don't have a unique value to substitute, warn and
# leave the entry unchanged.
# (We don't exit right away so that we can detect multiple problems
# within this genbki.pl run.)
sub lookup_oids
{
	my ($lookup, $catname, $attname, $lookup_opt, $bki_values, @lookupnames)
	  = @_;

	my @lookupoids;
	foreach my $lookupname (@lookupnames)
	{
		my $lookupoid = $lookup->{$lookupname};
		if (defined($lookupoid) and $lookupoid ne 'MULTIPLE')
		{
			push @lookupoids, $lookupoid;
		}
		else
		{
			push @lookupoids, $lookupname;
			if ($lookupname eq '-' or $lookupname eq '0')
			{
				if (!$lookup_opt)
				{
					warn sprintf
					  "invalid zero OID reference in %s.dat field %s line %s\n",
					  $catname, $attname, $bki_values->{line_number};
					$num_errors++;
				}
			}
			else
			{
				warn sprintf
				  "unresolved OID reference \"%s\" in %s.dat field %s line %s\n",
				  $lookupname, $catname, $attname, $bki_values->{line_number};
				$num_errors++;
			}
		}
	}
	return @lookupoids;
}

# ....

And the error message reflects that too:

unresolved OID reference "color_out" in pg_type.dat field typoutput line

“Wait. How is my function not accessible to the Perl script?”

After some debugging, I found out that there was a Meson config in the ADT directory, and it was missing my ADT’s source file, so I added it:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/meson.build;h=8515cd93650492cb79c74f58a5123a4ca59d9e31;hb=HEAD)

# Copyright (c) 2022-2023, PostgreSQL Global Development Group

backend_sources += files(
  'acl.c',
  'amutils.c',
  'array_expanded.c',
  'array_selfuncs.c',
  'array_typanalyze.c',
  'array_userfuncs.c',
  'arrayfuncs.c',
  'arraysubs.c',
  'arrayutils.c',
  'ascii.c',
  'bool.c',
  'cash.c',
  'char.c',
  'color.c', <-- added my ADT
  'cryptohashfuncs.c',
  'date.c',
  'datetime.c',
  ....

I also needed to modify the Makefile inside src/backend/utils/adt to include the object file for my ADT:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/Makefile;h=0de0bbb1b8ac7ccb60ea1ab39dbadf7004ea7586;hb=HEAD)

#
# Makefile for utils/adt
#
# src/backend/utils/adt/Makefile
#

subdir = src/backend/utils/adt
top_builddir = ../../../..
include $(top_builddir)/src/Makefile.global

override CPPFLAGS := -I. -I$(srcdir) $(CPPFLAGS)

# keep this list arranged alphabetically or it gets to be a mess
OBJS = \
	acl.o \
	amutils.o \
	array_expanded.o \
	array_selfuncs.o \
	array_typanalyze.o \
	array_userfuncs.o \
	arrayfuncs.o \
	arraysubs.o \
	arrayutils.o \
	ascii.o \
	bool.o \
	cash.o \
	char.o \
	color.o \  <-- added my ADT
	cryptohashfuncs.o \
	....

Now, on trying again, the compilation process worked perfectly. Time to test it out!

Edit: If you change the system catalog (.dat files), you would need to re-initalize the DB:

$ rm -rf test-data
$ ./test-install/bin/initdb $(pwd)/test-data # Re-create the database

Initial test with the color type ¶

I created a table with my color type:

postgres=# CREATE TABLE color (val color);
CREATE TABLE
postgres=# INSERT INTO color VALUES ('FF0000');
INSERT 0 1
postgres=# SELECT * FROM color;
  val
--------
 FF0000
(1 row)

postgres=# \d color
              Table "public.color"
 Column | Type  | Collation | Nullable | Default
--------+-------+-----------+----------+---------
 val    | color |           |          |

A sigh of relief. It worked!

This was a good start, but the type is of no use if we can’t perform operations on top of it. For getting started, I wanted to go ahead with an operator that “mixes” two colors and returns the result of type color.

Implementing operators for the color type ¶

Color addition is tricky, there are different techniques to perform it; but to keep things simple, I chose additive mixing as the algorithm for color addition. I was able to implement this by taking help from some articles I found online:

/*
 * color_add - adds two values of type color and returns the result
 */
PG_FUNCTION_INFO_V1(color_add);
Datum color_add(PG_FUNCTION_ARGS)
{
    color arg1 = PG_GETARG_UINT32(0);
    color arg2 = PG_GETARG_UINT32(1);

    // Extract the Red (R) channel component of arg1 and arg2 by shifting the bits
    // to the right and applying a bit mask to isolate the 8 bits representing the Red channel.
    uint32 r1 = (arg1 >> COLOR_RED_SHIFT) & COLOR_CHANNEL_MASK;
    uint32 r2 = (arg2 >> COLOR_RED_SHIFT) & COLOR_CHANNEL_MASK;

    // Extract the Green (G) channel component.
    uint32 g1 = (arg1 >> COLOR_GREEN_SHIFT) & COLOR_CHANNEL_MASK;
    uint32 g2 = (arg2 >> COLOR_GREEN_SHIFT) & COLOR_CHANNEL_MASK;

    // Extract the Blue (B) channel component.
    uint32 b1 = (arg1 >> COLOR_BLUE_SHIFT) & COLOR_CHANNEL_MASK;
    uint32 b2 = (arg2 >> COLOR_BLUE_SHIFT) & COLOR_CHANNEL_MASK;

    // Add the Red (R) channel components, ensuring that the result does not exceed
    // the maximum value represented by COLOR_CHANNEL_MASK (0xFF), which corresponds
    // to the upper limit of the Red channel. This prevents overflow.
    uint32 r = Min(r1 + r2, COLOR_CHANNEL_MASK);

    // Add the Green (G) channel components.
    uint32 g = Min(g1 + g2, COLOR_CHANNEL_MASK);

    // Add the Blue (B) channel components.
    uint32 b = Min(b1 + b2, COLOR_CHANNEL_MASK);

    // Reconstruct the color by left-shifting the Red (R), Green (G), and Blue (B)
    // components back to their original bit positions and combining them with the OR (|) operation.
    color result_color = (r << COLOR_RED_SHIFT) | (g << COLOR_GREEN_SHIFT) | (b << COLOR_BLUE_SHIFT);

    // Return the final color as a uint32 value.
    PG_RETURN_UINT32(result_color);
}

In addition to this, I thought it would be a good idea to add support for equality and inequality operators as the implementation would be fairly easy, due to the internal representation of color being a uint32:

/*
 * color_eq - checks for equality between two values of type color
 */
PG_FUNCTION_INFO_V1(color_eq);
Datum color_eq(PG_FUNCTION_ARGS)
{
    color arg1 = PG_GETARG_UINT32(0);
    color arg2 = PG_GETARG_UINT32(1);

    PG_RETURN_BOOL(arg1 == arg2);
}

/*
 * color_ne - checks for inequality between two values of type color
 */
PG_FUNCTION_INFO_V1(color_ne);
Datum color_ne(PG_FUNCTION_ARGS)
{
    color arg1 = PG_GETARG_UINT32(0);
    color arg2 = PG_GETARG_UINT32(1);

    PG_RETURN_BOOL(arg1 != arg2);
}

Nice! Now, I had support for these operators: =, <> and +.

To register these functions in the catalog (pg_proc.dat), I followed the same process:

{ oid => '7905', proname => 'color_eq', prorettype => 'bool',
  proargtypes => 'color color', prosrc => 'color_eq' },
{ oid => '7906', proname => 'color_ne', prorettype => 'bool',
  proargtypes => 'color color', prosrc => 'color_ne' },
{ oid => '7907', proname => 'color_add', prorettype => 'color',
  proargtypes => 'color color', prosrc => 'color_add' },

In order to register operators for the color type, I added the following to the pg_operator system catalog (pg_operator.dat):

# custom color type

{ oid => '7915', descr => 'equal',
  oprname => '=', oprcanmerge => 't', oprcanhash => 't', oprleft => 'color',
  oprright => 'color', oprresult => 'bool', oprcom => '=(color,color)',
  oprnegate => '<>(color,color)', oprcode => 'color_eq' },
{ oid => '7916', descr => 'not equal',
  oprname => '<>', oprcanmerge => 't', oprcanhash => 't', oprleft => 'color',
  oprright => 'color', oprresult => 'bool', oprcom => '<>(color,color)',
  oprnegate => '=(color,color)', oprcode => 'color_ne' },
{ oid => '7917', descr => 'add',
  oprname => '+', oprleft => 'color', oprright => 'color', oprresult => 'color',
  oprcom => '+(color,color)', oprcode => 'color_add' },

Here’s what these fields mean:

oid: OID (Object Identifier)
descr: A short description of the operator.
oprname: The name of the operator. It is the symbol or name used to invoke the operator, for example: =, <>, or +.
oprcanmerge: This is a boolean flag (t/f) indicating whether the operator supports merge joins. (I’ve kept it as t for now, but I’d be glad to know the suitable value for this.)
oprcanhash: This is a boolean flag (t/f) indicating whether the operator supports hash joins. (I’ve kept it as t for now.)
oprleft: The left operand data type of the operator.
oprright: The right operand data type of the operator.
oprresult: The data type of the result produced by the operator.
oprcom: This specifies the commutator operator. The commutator is another operator that has the same functionality but with the operand order reversed. For example, = and = are commutator operators of each other.
oprnegate: The negate operator for a comparison operator. It represents the opposite of the operator. For =, its negate is <>.
oprcode: This refers to the function implementing the operator.

The final test ¶

To test it out, I recompiled and was excited to see if the operators worked perfectly:

postgres=# CREATE TABLE color (val color);
CREATE TABLE
postgres=# INSERT INTO color VALUES ('FF0000');
INSERT 0 1
postgres=# SELECT * FROM color;
  val
--------
 FF0000
(1 row)

postgres=# SELECT val+'00FF00' FROM color;
 ?column?
----------
 FFFF00
(1 row)

postgres=# SELECT 'FF0000'::color + '00FF00'::color;
 ?column?
----------
 FFFF00
(1 row)


postgres=# SELECT * FROM color WHERE val='FF0000';
  val
--------
 FF0000
(1 row)

postgres=# SELECT * FROM color WHERE val<>'FF0000';
 val
-----
(0 rows)

Yay! The operators work as intended, and I was finally relieved to see how it turned out. It may not be perfect, but it was exciting to work on!

I’ve pushed these changes here, please feel free to check it out!

Conclusion ¶

Yes, I should have RTFM’d it. But honestly, this was more fun. If I had already knew about which routines to use, how to structure things, it wouldn’t have been interesting.

I spent 2 days on this, and I’d like to say: it was worth it. I’ll take a break, enjoy a few games and come back to hacking Postgres once again (yes, I have something planned).

Please email (contact <at> databases.systems) or DM on Twitter if you have any questions, corrections, or if you are hiring.

pg_experiments, Part 1: Modifying operator logic

2023-10-30T00:00:00Z

Exploring Postgres internals by changing operator logic.

Disclaimer: This is my first time diving deep into Postgres source code. I wrote this post after my short, 1-hour exploration in the code base, so I'm unaware of any ideal tricks/techniques to use for implementing things.

~~This is my first microblog! I’ll be writing short-form content for smaller projects/experiments.~~ Edit: Moved this to the main blog.

Today’s experiment was something which I thought would be fun to write. Thanks to Phil Eaton for helping me out with this experiment!

Why? ¶

During a conversation, Phil wrote this:

I’ve gone looking for where builtin SQL functions like +, - are defined and got lost. So if someone wants to find that and share that would be nice! 🙂

I got curious and wanted to know more. And since the job market has been meh this year, I decided to explore this further to overcome the monotonic, boring job search.

The goal for this exploration is to find where this logic is present, and tweak it to make the + operator subtract the passed values.

Diving deep into source code ¶

Since Phil mentioned about the + operator, I thought it would be a good start to start checking references for “add” as a keyword for the search.

I came across numeric.c in src/backend/utils/adt, which happens to be a gigantic ~10k+ LOC file. On doing a few more searches, I found something which does addition: numeric_add:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l2842

/*
 * numeric_add() -
 *
 *	Add two numerics
 */
Datum
numeric_add(PG_FUNCTION_ARGS)
{
	Numeric		num1 = PG_GETARG_NUMERIC(0);
	Numeric		num2 = PG_GETARG_NUMERIC(1);
	Numeric		res;

	res = numeric_add_opt_error(num1, num2, NULL);

	PG_RETURN_NUMERIC(res);
}

/*
 * numeric_add_opt_error() -
 *
 *	Internal version of numeric_add().  If "*have_error" flag is provided,
 *	on error it's set to true, NULL returned.  This is helpful when caller
 *	need to handle errors by itself.
 */
Numeric
numeric_add_opt_error(Numeric num1, Numeric num2, bool *have_error)
{
	NumericVar	arg1;
	NumericVar	arg2;
	NumericVar	result;
	Numeric		res;

	/*
	 * Handle NaN and infinities
	 */
	if (NUMERIC_IS_SPECIAL(num1) || NUMERIC_IS_SPECIAL(num2))
	{
		if (NUMERIC_IS_NAN(num1) || NUMERIC_IS_NAN(num2))
			return make_result(&const_nan);
		if (NUMERIC_IS_PINF(num1))
		{
			if (NUMERIC_IS_NINF(num2))
				return make_result(&const_nan); /* Inf + -Inf */
			else
				return make_result(&const_pinf);
		}
		if (NUMERIC_IS_NINF(num1))
		{
			if (NUMERIC_IS_PINF(num2))
				return make_result(&const_nan); /* -Inf + Inf */
			else
				return make_result(&const_ninf);
		}
		/* by here, num1 must be finite, so num2 is not */
		if (NUMERIC_IS_PINF(num2))
			return make_result(&const_pinf);
		Assert(NUMERIC_IS_NINF(num2));
		return make_result(&const_ninf);
	}

	/*
	 * Unpack the values, let add_var() compute the result and return it.
	 */
	init_var_from_num(num1, &arg1);
	init_var_from_num(num2, &arg2);

	init_var(&result);
	add_var(&arg1, &arg2, &result);

	res = make_result_opt_error(&result, have_error);

	free_var(&result);

	return res;
}

numeric_add_opt_error calls add_var, which has sign handling logic, and calls other functions for addition/subtraction based on signs:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l8353

/*
 * add_var() -
 *
 *	Full version of add functionality on variable level (handling signs).
 *	result might point to one of the operands too without danger.
 */
static void
add_var(const NumericVar *var1, const NumericVar *var2, NumericVar *result)
{
	/*
	 * Decide on the signs of the two variables what to do
	 */
	if (var1->sign == NUMERIC_POS)
	{
		if (var2->sign == NUMERIC_POS)
		{
			/*
			 * Both are positive result = +(ABS(var1) + ABS(var2))
			 */
			add_abs(var1, var2, result);
			result->sign = NUMERIC_POS;
		}
		else
		{
			/*
			 * var1 is positive, var2 is negative Must compare absolute values
			 */
			switch (cmp_abs(var1, var2))
			{
				case 0:
					/* ----------
					 * ABS(var1) == ABS(var2)
					 * result = ZERO
					 * ----------
					 */
					zero_var(result);
					result->dscale = Max(var1->dscale, var2->dscale);
					break;

				case 1:
					/* ----------
					 * ABS(var1) > ABS(var2)
					 * result = +(ABS(var1) - ABS(var2))
					 * ----------
					 */
					sub_abs(var1, var2, result);
					result->sign = NUMERIC_POS;
					break;

				case -1:
					/* ----------
					 * ABS(var1) < ABS(var2)
					 * result = -(ABS(var2) - ABS(var1))
					 * ----------
					 */
					sub_abs(var2, var1, result);
					result->sign = NUMERIC_NEG;
					break;
			}
		}
	}
	else
	{
		if (var2->sign == NUMERIC_POS)
		{
			/* ----------
			 * var1 is negative, var2 is positive
			 * Must compare absolute values
			 * ----------
			 */
			switch (cmp_abs(var1, var2))
			{
				case 0:
					/* ----------
					 * ABS(var1) == ABS(var2)
					 * result = ZERO
					 * ----------
					 */
					zero_var(result);
					result->dscale = Max(var1->dscale, var2->dscale);
					break;

				case 1:
					/* ----------
					 * ABS(var1) > ABS(var2)
					 * result = -(ABS(var1) - ABS(var2))
					 * ----------
					 */
					sub_abs(var1, var2, result);
					result->sign = NUMERIC_NEG;
					break;

				case -1:
					/* ----------
					 * ABS(var1) < ABS(var2)
					 * result = +(ABS(var2) - ABS(var1))
					 * ----------
					 */
					sub_abs(var2, var1, result);
					result->sign = NUMERIC_POS;
					break;
			}
		}
		else
		{
			/* ----------
			 * Both are negative
			 * result = -(ABS(var1) + ABS(var2))
			 * ----------
			 */
			add_abs(var1, var2, result);
			result->sign = NUMERIC_NEG;
		}
	}
}

As we saw above, add_var uses the add_abs and sub_abs functions based on the signs. I took a look at add_abs, and it seems to be the exact place where the actual logic happens (carry bits, etc.):

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l11331

/*
 * add_abs() -
 *
 *	Add the absolute values of two variables into result.
 *	result might point to one of the operands without danger.
 */
static void
add_abs(const NumericVar *var1, const NumericVar *var2, NumericVar *result)
{
	NumericDigit *res_buf;
	NumericDigit *res_digits;
	int			res_ndigits;
	int			res_weight;
	int			res_rscale,
				rscale1,
				rscale2;
	int			res_dscale;
	int			i,
				i1,
				i2;
	int			carry = 0;

	/* copy these values into local vars for speed in inner loop */
	int			var1ndigits = var1->ndigits;
	int			var2ndigits = var2->ndigits;
	NumericDigit *var1digits = var1->digits;
	NumericDigit *var2digits = var2->digits;

	res_weight = Max(var1->weight, var2->weight) + 1;

	res_dscale = Max(var1->dscale, var2->dscale);

	/* Note: here we are figuring rscale in base-NBASE digits */
	rscale1 = var1->ndigits - var1->weight - 1;
	rscale2 = var2->ndigits - var2->weight - 1;
	res_rscale = Max(rscale1, rscale2);

	res_ndigits = res_rscale + res_weight + 1;
	if (res_ndigits <= 0)
		res_ndigits = 1;

	res_buf = digitbuf_alloc(res_ndigits + 1);
	res_buf[0] = 0;				/* spare digit for later rounding */
	res_digits = res_buf + 1;

	i1 = res_rscale + var1->weight + 1;
	i2 = res_rscale + var2->weight + 1;
	for (i = res_ndigits - 1; i >= 0; i--)
	{
		i1--;
		i2--;
		if (i1 >= 0 && i1 < var1ndigits)
			carry += var1digits[i1];
		if (i2 >= 0 && i2 < var2ndigits)
			carry += var2digits[i2];

		if (carry >= NBASE)
		{
			res_digits[i] = carry - NBASE;
			carry = 1;
		}
		else
		{
			res_digits[i] = carry;
			carry = 0;
		}
	}

	Assert(carry == 0);			/* else we failed to allow for carry out */

	digitbuf_free(result->buf);
	result->ndigits = res_ndigits;
	result->buf = res_buf;
	result->digits = res_digits;
	result->weight = res_weight;
	result->dscale = res_dscale;

	/* Remove leading/trailing zeroes */
	strip_var(result);
}

A simple explanation of how it works:

add_abs function adds the absolute values of two NumericVars, var1 and var2, and stores the result in result.
A loop iterates through the digits of result, adding digits from var1 and var2 at corresponding positions while accounting for any carry from the previous iteration.
If the sum of digits exceeds NBASE, the carry is propagated to the next digit, and the result digit is set accordingly.
After the loop, an assertion checks that no carry remains to ensure a successful addition.
The function updates the result and removes leading/trailing zeroes to maintain a canonical form.

The result is stored in a NumericVar type, which has a few fields to store the values and some metadata (signs, etc.). The documentation is well-written:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l254

typedef struct NumericVar
{
	int			ndigits;		/* # of digits in digits[] - can be 0! */
	int			weight;			/* weight of first digit */
	int			sign;			/* NUMERIC_POS, _NEG, _NAN, _PINF, or _NINF */
	int			dscale;			/* display scale */
	NumericDigit *buf;			/* start of palloc'd space for digits[] */
	NumericDigit *digits;		/* base-NBASE digits */
} NumericVar;

Here’s an interesting comment talking about the format and NBASE:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l55)

 * Numeric values are represented in a base-NBASE floating point format.
 * Each "digit" ranges from 0 to NBASE-1.  The type NumericDigit is signed
 * and wide enough to store a digit.  We assume that NBASE*NBASE can fit in
 * an int.  Although the purely calculational routines could handle any even
 * NBASE that's less than sqrt(INT_MAX), in practice we are only interested
 * in NBASE a power of ten, so that I/O conversions and decimal rounding
 * are easy.  Also, it's actually more efficient if NBASE is rather less than
 * sqrt(INT_MAX), so that there is "headroom" for mul_var and div_var_fast to
 * postpone processing carries.
 *
 * Values of NBASE other than 10000 are considered of historical interest only
 * and are no longer supported in any sense; no mechanism exists for the client
 * to discover the base, so every client supporting binary mode expects the
 * base-10000 format.  If you plan to change this, also note the numeric
 * abbreviation code, which assumes NBASE=10000.
 * ----------
 */

Making changes ¶

Even before making any changes, I had to set up a working environment to compile and test changes. Thankfully, Phil had posted this set of commands with comments for getting started with a simple setup:

$ ./configure --prefix=$(pwd)/test-install
$ make -j8 && make install
$ ./test-install/bin/initdb $(pwd)/test-data # Create a database
$ echo "port=8080" > $(pwd)/postgres.conf # This is the minimal postgres run config
$ mkdir test-run # Place for postgres to put lock files
$ ./test-install/bin/postgres --config-file=$(pwd)/postgres.conf -D $(pwd)/test-data -k $(pwd)/test-run

Now that I had some context on how numeric addition works in Postgres, making the change was really easy; all I had to do was change the call from add_var to sub_var in numeric_add_opt_error. (keep in mind that I’m ignoring other cases)

Diff:

diff --git a/src/backend/utils/adt/numeric.c b/src/backend/utils/adt/numeric.c
index bf61fd7dbc..e174f9a9ff 100644
--- a/src/backend/utils/adt/numeric.c
+++ b/src/backend/utils/adt/numeric.c
@@ -2905,7 +2905,7 @@ numeric_add_opt_error(Numeric num1, Numeric num2, bool *have_error)
        init_var_from_num(num2, &arg2);

        init_var(&result);
-       add_var(&arg1, &arg2, &result);
+       sub_var(&arg1, &arg2, &result);

        res = make_result_opt_error(&result, have_error);

Let’s recompile and restart the server:

$ make -j8 && make install
$ ./test-install/bin/postgres --config-file=$(pwd)/postgres.conf -D $(pwd)/test-data -k $(pwd)/test-run

And run the shell:

$ ./test-install/bin/psql -p 8080 -h localhost -d postgres

To test it out, I created a new table (to prevent possible constant folding and just to be extra sure that it works):

postgres=# CREATE TABLE numbers (num1 numeric, num2 numeric);
CREATE TABLE
postgres=# INSERT INTO numbers (num1, num2) VALUES (5.25, 3.75);
INSERT 0 1

And the moment I had been waiting for:

postgres=# SELECT num1 + num2 AS sum FROM numbers;
 sum
------
 1.50
(1 row)

Yay! 5.25 - 3.75 = 1.50, so it’s working as intended!

What’s next? ¶

The above exploration was something I managed to do in an hour; I’m currently looking to explore more and have a project planned where I’d try to exploit Postgres internals and build something out of it!

Please email (contact <at> databases.systems) or DM on Twitter if you have any questions, corrections, or if you are hiring.

An attempt at visualizing the Go GC

2023-10-09T00:00:00Z

A rejected talk proposal for GopherCon India '23.

Disclaimer: I have limited experience with GCs and the Go runtime, so if there are any mistakes, please go easy on me! This post doesn't focus on understanding the Go GC; those who'd like to do so should read the source code once. I'll cover them in a future post if possible.

In late 2022, I stumbled upon this visualization during my exploration around the Go GC (for learning the internals). A thought came to my mind, “Wouldn’t it be pretty cool if someone could build a tool to visualize the live process of mark and sweep?”, and I got more curious than ever.

This was more tempting than anything else I had ever worked on. Why? Because I always wanted to explore Go internals and what better way to explore than to hack around stuff and find out more?

I talked to a few people, explained about the experiment, and most of them wanted me to publish this post, irrespective of how much progress I made on the experiment, so, here I am.

(My original 2022 playground got nuked due to a boot failure after a Windows Update, and I eventually had to reset the entire system. (thank you Windows!) ¹)

What was I trying to build? ¶

As explained above, Pusher’s GC visualizer was the starting point of this experiment. Their visualizer was meant for educational purposes. “What would it take to build a visualizer with real data?” was the real question.

Okay, wait. But even before all of this, is it even required? Probably not. I don’t know if a standalone visualizer would help in actual production use cases. But in true hacker spirit, I had to do it.

The end goal for this exploration is to build a tool that:

attempts to extract maximum information from the Go runtime and treats it as the source of truth
Use the extracted information to build relationships between individual runtime “objects”
Find a way to determine object liveness from the extracted info, or different sources
Construct an object graph with object liveness info to color “objects” to emulate the tri-color mark-sweep algorithm
Iterate over different object graphs and generate an animated overview of the GC process using real data

Simply put, I wanted to build a tool which would show a “live” object graph with objects colored as black/grey/white, with the colors changing with respect to time.

To understand things further, I had to learn some GC theory. Let me give you a walkthrough about a few things in GC theory.

Understanding some GC theory ¶

Okay, first things first. What is the heap?

In simple terms, a heap is either a continuous array of memory words or a non-contiguous collection of memory blocks, where each word is a sequence of bytes large enough to hold a pointer. Allocating objects on the heap is usually meant for objects whose size, quantity, or lifetime cannot be determined at compile-time. Think of it as a “free store” that can be managed manually or automatically. The heap is usually represented as an object graph, where the nodes are heap objects and the edges are references to the heap objects.

In manual memory management, if we incorrectly use explicit deallocation to reclaim memory, we may accidentally free memory that has references to it. These references are often called as dangling pointers. Following dangling pointers may result in unpredictable behavior.

This is where GC (garbage collection) comes into play and helps in preventing the creation of dangling pointers. The garbage collector attempts to reclaim memory that was allocated by the program, but is no longer referenced. The GC has the context of the global heap state, which helps in making reclamation decisions easier.

Remember, GC is not a silver bullet, and manual memory management can be performed correctly with care. I do not want to start a war :)

Next are mutators and collectors, the two most commonly discussed “actors” in the context of garbage collection.

The mutator allocates new objects and “mutates” the (heap) object graph by changing references. A collector discovers unreachable objects (objects that cannot be reached by following any set of edges from the roots).

Bear with me, some more theory incoming:

An object is said to be live if it will be accessed sometime in the near future by the mutator. However, figuring out object liveness is a hard problem: liveness is non-deterministic; there is technically no possible way to determine if an object will be accessed by the mutator before time X.

A way to figure out if an object is live is to approximate it. Approximation can be done by object reachability. (check if an object X is reachable from an object Y by following references from Y -> X)

The algorithm that Go’s garbage collector uses is the mark and sweep algorithm. The mark and sweep algorithm is a two-phase garbage collection method that first marks all live objects in memory and then sweeps through the memory to deallocate unreferenced objects. In mark-and-sweep, the collector operates in two phases:

Tracing (mark) phase: The collector traverses through the object graph by starting from the roots, and then marks each object it finds during its traversal.
Sweeping phase: Finds unmarked objects, and determines them as garbage indirectly (mark-and-sweep doesn’t directly detect garbage, but assumes that unmarked objects are garbage.)

Go’s GC uses the concept of tri-color abstraction to efficiently track and identify live objects for garbage collection without stopping program execution for significant periods of time. (this is achieved by marking objects during mutation)

Under the tri-color abstraction, the object graph is partitioned into three colors (black, white, and grey). Initially, every node is white and the root is grey (as it is the first node that is encountered).

When an object is encountered by the collector, it is marked as grey, and when the object is scanned and its descendants are identified it is marked as black. Gray objects are those that are potentially reachable and need further examination, while black objects are confirmed as live and in use. Objects that remain white after this process can be safely deallocated.

Mark-and-sweep also focuses on preserving an invariant to guarantee correctness: at the end of each iteration of the marking loop, there are no references from black to white objects. White objects must be only reachable from a grey object.

What would happen if the invariant is to be broken? Let’s consider a black object pointing to a white object after the mark phase. As the algorithm doesn’t process black objects further, the white object will be reclaimed by the collector during the sweep phase, leading to the presence of dangling pointers.

Phew! Hope I didn’t bore you out this wall of text, but this would help understand the content below in a better way.

Exploring the runtime package and more ¶

Everything we’re looking at right now has some overlap with the runtime. I knew that Go provides a runtime package, but I wasn’t sure if it exported any information related to the Go GC.

After some exploration, the closest related thing that I could find in the runtime package was runtime.MemStats. But MemStats exports numerical data; my goal isn’t to build a live chart of memory statistics, a lot of great tools do it already. I needed more information, numerical data isn’t going to help in any way.

Thankfully during my search, I came across Ardan Labs’ article on GC semantics, where a section on GC traces piqued my interest. GC traces seemed promising as they could be toggled using a flag and had a low entry barrier. However, the trace itself doesn’t have sufficient information on actual objects and I don’t blame it; traces aren’t meant to capture that information anyway.

I had to look for other solutions.

Forking Go ¶

Yep, you read it right. I was tempted, really tempted. Let me explain this journey of forking Go.

So, nothing from the official Go standard library had been of help so far.

If I were to get the most information, the only place to look into would be the actual Go source code. runtime/mgc.go is the main GC algorithm, and is the expected place where I should be making tweaks to export data related to the mark-and-sweep process. (or maybe in runtime/mgcmark.go?)

But on further exploration and lots of code reading, I just couldn’t find a way to cleanly export this data. Every patch I tried to add resulted in more tweaks being made to different parts of the runtime. ²

The major challenge/disadvantage of using this approach was that I needed to maintain a forked version of Go. This also made the exploration more divergent than my intended goal - the idea was to make it run on the official Go releases, not forked versions.

Finally, I decided to stop making progress on this idea.

Heap dumps ¶

At this point, I was lost. I did some more research and found a Go proposal which was interesting and close to my experiment’s goals.

After reading the title, I thought - “wow. why didn’t I ever think of using heap dumps?”

Heap dumps are snapshots of the program’s memory at a specific point in time.

The proposal talks about the various ways to obtain a heap dump: by using debug.WriteHeapDump, using standard core dumps, etc.

Since I wanted to stick closer to the Go runtime, I decided to go ahead with the heap dumps generated by Go as the source for obtaining more information.

I went straight to the playground and generated a heap dump in a few seconds. But I still had to parse it, and oh boy, I didn’t know that a working heap dump parser for Go 1.18 (the version I was working on at that time (Oct 2022)) didn’t exist.

Keith Randall’s heapdump14 was the standard heap dump parser for Go, but it wasn’t maintained for newer versions. Keeping up with the runtime is a challenging task.

Keith also mentioned that the GC has been exporting less information related to types in the recent versions. (I need you to keep this in mind for the sections below)

I looked up more parsers, but each parser had a few bugs that prevented me from getting a complete overview of the heap dump. ³

Understanding the heap dump format ¶

Now, I had to build a simple heap dump parser on my own. To start working on the implementation, I needed to understand the heap dump format.

The Go wiki nicely documents this. A Go heap dump is essentially a sequence of records, where each record starts with a 64-bit unsigned integer describing the type of the record.

There are 17 record types:

0 = EOF
1 = object
2 = otherroot
3 = type
4 = goroutine
5 = stack frame
6 = dump params
7 = registered finalizer
8 = itab
9 = OS thread
10 = mem stats
11 = queued finalizer
12 = data segment
13 = bss segment
14 = defer record
15 = panic record
16 = alloc/free profile record
17 = alloc stack trace sample

To know more about how Go writes these heap dumps, I’d recommend checking the bonus section at the end of this post.

Parse, parse, parse ¶

After understanding the heap dump format, I started working on building the parser. The basic approach for parsing it was pretty simple: identify the record type using the 64-bit unsigned integer, use a reader for that record type, and return the read contents.

The actual implementation looks something like this:

// parse.go

func ParseDump(rd *bufio.Reader) ([]RecordData, error) {
    err := record.ReadHeader(rd)
    if err != nil {
        return nil, err
    }

    ....

    for {
        r, err := record.ReadRecord(rd)
        if err != nil {
            return nil, err
        }
        
        ....

ReadRecord is the place where the type of record is identified and a suitable reader is initialized:

// record.go

func ReadRecord(r *bufio.Reader) (Record, error) {
    // Read record type
    recordType, err := binary.ReadUvarint(r)
    if err != nil {
        return nil, err
    }

    var record Record

    switch RecordType(recordType) {
    case EOF:
        record = &EOFRecord{}
    case Object:
        record = &ObjectRecord{}
    case OtherRoot:
        record = &OtherRootRecord{}
    case TypeDescriptor:
        record = &TypeDescriptorRecord{}
    case Goroutine:
        record = &GoroutineRecord{}
    case StackFrame:
        record = &StackFrameRecord{}
    case DumpParams:
        record = &DumpParamsRecord{}
    case RegisteredFinalizer:
        record = &RegisteredFinalizerRecord{}
    case ITab:
        record = &ITabRecord{}
    case OSThread:
        record = &OSThreadRecord{}
    case MemStats:
        record = &MemStatsRecord{}
    case QueuedFinalizer:
        record = &QueuedFinalizerRecord{}
    case DataSegment:
        record = &DataSegmentRecord{}
    case BSSSegment:
        record = &BSSSegmentRecord{}
    case DeferRecord:
        record = &DeferRecordRecord{}
    case PanicRecord:
        record = &PanicRecordRecord{}
    case AllocFreeProfile:
        record = &AllocFreeProfileRecord{}
    case AllocStackTraceSample:
        record = &AllocStackTraceSampleRecord{}
    default:
        return nil, fmt.Errorf("unexpected record type: %v", recordType)
    }

    // Read record. Each record implements the Record interface.
    err = record.Read(r)
    if err != nil {
        return nil, err
    }

    return record, nil
}

For example, if the record type is itab, then the reader’s implementation will look like this:

// itab.go

// uvarint: Itab address
// uvarint: address of type descriptor for contained type
type ITabRecord struct {
    Address            uint64
    TypeDescriptorAddr uint64
}

func (r *ITabRecord) GetAddress() uint64 {
    return r.Address
}

func (r *ITabRecord) Read(rd *bufio.Reader) error {
    var err error

    r.Address, err = binary.ReadUvarint(rd)
    if err != nil {
        return err
    }

    r.TypeDescriptorAddr, err = binary.ReadUvarint(rd)
    if err != nil {
        return err
    }

    return nil
}

By using this approach, I can get a textual representation of the entire heap:

OS thread at address 0x54b300 (Go internal ID = 0, OS ID = 53413)
Data segment at address 0x546660 (content size = 16784, pointers = 1184)
BSS segment at address 0x54a800 (content size = 196376, pointers = 10185)
Registered finalizer at address 0xc000078060, FuncVal ptr address = 0x4b7078, type = 4832448, object type = 4832448
Registered finalizer at address 0xc0000a8000, FuncVal ptr address = 0x4b7078, type = 4832448, object type = 4832448
Registered finalizer at address 0xc0000a8060, FuncVal ptr address = 0x4b7078, type = 4832448, object type = 4832448
Registered finalizer at address 0xc0000a80c0, FuncVal ptr address = 0x4b7078, type = 4832448, object type = 4832448

While I reached near completion on this parser, I got to know that a similar parser existed in the heapspurs tool, which I wasn’t aware of. (Sigh. Why didn’t I come across it earlier?)

heapspurs tries to do a lot more, but the core idea of parsing remains the same - identify record types using the 64-bit unsigned integer and read the contents based on the record type.

However, I wanted the result to be somewhat similar to Eclipse Jifa’s heap dump view. So, I separated the parsing logic into a different project (heaputil) (in case anyone else wants to use it) and decided to build a separate project (heapview) for the visualizer.

Jifa's view

I built the records view using HTML templates:

<!-- Records view -->
    <div id="recordsView">
        {{range .Records}}
            <div class="row {{.RecordType}}">
                {{.Repr}}
                {{if .HasPointers}}
                    <button onclick="toggleRow('{{.RowID}}')">Toggle</button>
                    <div id="{{.RowID}}" class="pointer-info" style="display: none;">
                        {{range .Pointers}}
                            <p>Pointer({{.Index}}) at address 0x{{.Address}} (incoming = 0x{{.Incoming}}, outgoing = 0x{{.Outgoing}})</p>
                        {{end}}
                    </div>
                {{end}}
            </div>
        {{end}}
    </div>

....

    data := templateData{
        RecordTypes:     GetUniqueRecordTypes(records),
        Records:         records,
        GraphVizContent: graphContent,
    }

    var htmlBuilder strings.Builder
    err = tmpl.Execute(&htmlBuilder, data)
    if err != nil {
        return "", err
    }

....

Which results in a fairly nice view of individual records:

Building object graphs ¶

The main part of the visualization is the object graph. An object graph is essentially a visual or conceptual map that shows how different objects in the heap are connected to each other through references or pointers.

Even before building the object graph, I had to decide on what record types are going to be a part of it. A small heap dump can have a lot of records, and representing everything in the object graph is not going to work out. (Note: I’m not talking about other types like Gs and Ms ⁴, but other “objects” like data segments and BSS segments.)

For the initial implementation, I wanted to keep the scope small - so, the focus was only on objects. No data or BSS segments.

What would a barebones implementation for this look like? After some head-scratching, I decided to do some additional things while parsing each record:

if the record isn’t of type Object, just ignore. (to keep the graph small)
create a node with a unique ID, where the label would have details (name, address, etc.)
the name and address is extracted from the record’s representation (Repr) (yup, a tiny hack)
create edges for outgoing pointers

    for {
        r, err := record.ReadRecord(rd)
        if err != nil {
            return dotContent.String(), err
        }

        _, isEOF := r.(*record.EOFRecord)
        if isEOF {
            break
        }

        dp, isDumpParams := r.(*record.DumpParamsRecord)
        if isDumpParams {
            dumpParams = dp
        }

        // Filter out objects. If the record isn't of the type Object, ignore.
        _, isObj := r.(*record.ObjectRecord)
        if !isObj {
            continue
        }

        // Create a DOT node for each record
        nodeName := fmt.Sprintf("Node%d", counter)
        counter++
        name, address := ParseNameAndAddress(r.Repr())
        nodeLabel := fmt.Sprintf("[%s] %s", name, address)

        // Write DOT node entry within the "heap" cluster
        s := fmt.Sprintf("    %s [label=\"%s\"];\n", nodeName, nodeLabel)
        dotContent.WriteString(s)

        // Check if the record has pointers
        p, isParent := r.(record.ParentGuard)
        if isParent {
            _, outgoing := record.ParsePointers(p, dumpParams)
            for i := 0; i < len(outgoing); i++ {
                if outgoing[i] != 0 {
                    childNodeName := fmt.Sprintf("Pointer0x%x", outgoing[i])

                    // Create an edge from the current record to the child record
                    s := fmt.Sprintf("    %s -> %s;\n", nodeName, childNodeName)
                    dotContent.WriteString(s)
                }
            }
        }
    }

Now, to render this to a HTML view, I used viz.js (ignore the terrible JS code, I’m not good at it):

<!-- Graph view (initially hidden) -->
    <div id="graphView" style="display: none;">
        <div id="graphContainer">
            <!-- The GraphViz graph will be rendered here using viz.js -->
        </div>
    </div>
</div>
<!-- .... -->
<script src="https://cdnjs.cloudflare.com/ajax/libs/viz.js/2.1.2/viz.js"></script>
<!-- .... -->
<script>
    // ....

    // Function to render the GraphViz graph
    function renderGraph(graphVizContent) {
        var graphContainer = document.getElementById("graphContainer");
        
        // Create a new Viz instance with the graph content and render it
        var viz = new Viz();
        viz.renderSVGElement(graphVizContent)
            .then(function(element) {
                // Append the SVG element to the container
                graphContainer.innerHTML = "";
                graphContainer.appendChild(element);
            })
            .catch(function(error) {
                console.error(error);
            });
    }
</script>

This finally came into life like this:

Source code is available here!

Coloring nodes ¶

With the object graph in place, all that was required to bring this experiment to an almost completed stage was to figure out the core “algorithm” to color individual nodes with any of these three colors (black, grey, white) to emulate the tri-color mark-sweep algorithm.

A rough sketch looked like this:

take two “snapshots” (heap dumps) at different times
compare the heap dumps and identify “dangling” objects
- a “dangling” object could be assumed to have no references; which means it is waiting to get GC’d (garbage collected)
- we can color these objects as white?
objects with existing pointers can be colored black, and their descendants can be colored grey.

This looked good enough, but I had to at least figure out the color coding “algorithm” on a single object graph; irrespective of temporal aspects. I could build on top on this foundational algorithm, and can add context in terms of time - seemed reasonable.

The foundational algorithm would be similar:

take a heap dump
construct the object graph
identify objects with pointers and color them black, and their descendants grey.
color objects without references as white

This assumes that objects which are “standalone” (without any references) are garbage and are waiting to be collected by the GC. The descendants of the black object are immediately colored grey to preserve the invariant which I mentioned above. (black objects cannot point to white objects)

Writing code for this part was a bit harder than I thought. I went ahead with parsing the original graph using regular expressions to distinguish between nodes and edges. If a node is encountered, we color it grey, and if has children nodes, we color it black. Here’s the hacky code I wrote for achieving this:

// NodeColor represents the color of a node in the graph.
type NodeColor int

const (
    White NodeColor = iota
    Grey
    Black
)

func ColorGraph(inputPath, outputPath string) error {
    file, err := os.Open(inputPath)
    if err != nil {
        return err
    }
    defer file.Close()

    nodeColors := make(map[string]NodeColor)

    // Parse the GraphViz file and determine node colors.
    err = determineNodeColors(file, nodeColors)
    if err != nil {
        return err
    }

    // Modify the graph to add node colors.
    err = applyColors(inputPath, outputPath, nodeColors)
    if err != nil {
        return err
    }

    return nil
}

func determineNodeColors(dotFile io.Reader, nodeColors map[string]NodeColor) error {
    scanner := bufio.NewScanner(dotFile)

    // Regular expression pattern to match node lines.
    nodePattern := regexp.MustCompile(`^\s+(\w+)\s+\[.*\]`)

    // Regular expression pattern to match edge lines.
    edgePattern := regexp.MustCompile(`^\s+(\w+)\s+->\s+(\w+)`)

    for scanner.Scan() {
        line := scanner.Text()
        if matches := nodePattern.FindStringSubmatch(line); matches != nil {
            // This is a node line.
            nodeName := matches[1]

            // Processed object, color it as grey.
            nodeColors[nodeName] = Grey
        } else if matches := edgePattern.FindStringSubmatch(line); matches != nil {
            // This is an edge line.
            sourceNode := matches[1]

            // Mark the source node as black.
            nodeColors[sourceNode] = Black
        }
    }

    return nil
}

func applyColors(inputPath, outputPath string, nodeColors map[string]NodeColor) error {
    content, err := os.ReadFile(inputPath)
    if err != nil {
        return err
    }

    // Create a regular expression pattern to match node lines.
    nodePattern := regexp.MustCompile(`^\s+(\w+)\s+\[.*\]`)

    lines := strings.Split(string(content), "\n")

    var modifiedContent strings.Builder

    for _, line := range lines {
        if matches := nodePattern.FindStringSubmatch(line); matches != nil {
            nodeName := matches[1]
            color, exists := nodeColors[nodeName]
            if exists {
                line = line[:len(line)-2]

                fontColor := "black"
                if color == Black {
                    fontColor = "white"
                }

                // Modify the line to include the fillcolor and fontcolor attributes with the determined color.
                line = fmt.Sprintf("%s fillcolor=%s, fontcolor=%s, style=filled];", line, getColorName(color), fontColor)
            }
        }
        modifiedContent.WriteString(line)
        modifiedContent.WriteString("\n")
    }

    // Write the modified GraphViz content to the output GraphViz file.
    err = os.WriteFile(outputPath, []byte(modifiedContent.String()), 0644)
    if err != nil {
        return err
    }

    return nil
}

This resulted in this object graph (trimmed for easier visualization):

What a relief. I finally had managed to figure out coloring individual nodes; and this brought me to closer to my end goal.

What’s next? ¶

I feel that more type information can be extracted from core dumps using viewcore. viewcore is a tool for exploring the state of a Go process that has dumped core. I had considered using viewcore, and even made a forum post for understanding more about it:

But since I was focusing on extracting information from heap dumps written by debug.WriteHeapDump, I couldn’t get much time to explore standard core dumps. Using core dumps might be a better way do it, as with more type information, the construction of object graphs becomes easier; and you get more detailed information as opposed to plain objects with addresses.

The animation part is easy - either iterate over a set of GraphViz graphs and generate a GIF, or parse each graph, check if the object references have changed or not and generate even smoother animations.

If you’d like to contribute or have better ideas, please create an issue on this repository. I’ll be happy to assist with any improvements!

I’ve open sourced the heap dump parser as heaputil, and the heap dump visualizer as heapview.

Conclusion ¶

And that is it! I never knew that I’d be spending my time on this, but it has been a great experience!

Please let me know if you have any comments or want to know more about the experiment! You can reach out to me via contact <at> databases.systems, Twitter DMs or LinkedIn.

I had backups configured for my main drive, but not for WSL (Windows Subsystem for Linux), where I had my entire playground. Windows is currently my primary operating system, and I’ve been enjoying some gaming on Steam during my free time. ↩︎
GC code has a lot of pointer logic which I don’t understand. Effectively patching to export information felt like a gargantuan task. ↩︎
gohat was one alternative to heapdump14, but it didn’t work. I also tried multiple forks of heapdump14, and gohat but they didn’t work as well. Someone also built a GNU poke pickle for reading Go heap dumps, but I wasn’t able to install poke after multiple tries, so I left it alone. ↩︎
A G is a goroutine, a M is an OS thread. The Go scheduler’s job is to match a G (code to execute), a M (where to execute), and a P (resources to execute it). ↩︎

Building a collaborative text editor in Go

2023-08-12T00:00:00Z

Exploring how pairpad was built in Go.

This post was originally written on 2022–12–17.

The first half of this article discusses about introducing real-time collaboration properties to text editors, and the second half discusses how pairpad was built.

Feel free to skip to the second part if you're already familiar with how collaborative editing works.

Introduction ¶

7 years ago, I built a text editor using Python 2.7 and Qt. It was one of my first projects, and it was fulfilling - seeing the editor’s buffer getting updated with the press of a key was a “wow” factor.

Recently, I wanted to try my hand at building text editors again, but I was looking for something more interesting. That’s when I decided to build a collaborative text editor called pairpad from scratch in Go.

“Why?”, you may ask - I thought it would be a fun challenge. And building it was exactly that - fun!

The primitives: text editors ¶

Before delving into collaborative text editors, let’s understand what a text editor actually is:

A text editor provides an interactive interface for editing text, typically supporting features such as cursor movement, text selection, copy and paste, and various text formatting options.

At the core of a text editor is typically a data structure that represents the text content. These data structures allow for dynamic and efficient insertion and deletion of characters without requiring extensive shifting of existing text.

When a user performs an operation on the text, such as insert or delete, the text editor translates these actions into low-level operations on the underlying data structure. For example, when a character is typed, the editor inserts it at the current cursor position by updating the data structure accordingly. Similarly, deleting a character involves removing it from the data structure and adjusting the positions of subsequent characters.

To support more complex operations, text editors often use additional data structures. These structures are used in implementing features like undo/redo, search and replace, syntax highlighting, code completion, and other advanced editing functionalities.

To understand what these operations look like, let’s take a look at this diagram:

Performing operations in a text editor

This diagram presents a basic depiction of how a text editor functions. In this example, the text editor has a buffer of size 4 (for demonstration purposes) and stores the content “WING” within the buffer.

An operation is performed to delete the character at position 3, resulting in the removal of the content “G” from that position.

Following that, an insertion operation is executed at the same position (3), adding the character “K”. The buffer now holds the content “WINK”.

Disclaimer: all editors might not work like this.

This is fine, but let’s look at how collaboration would take place in this setting.

Introducing collaboration in text editing ¶

Single-user text editing is relatively straightforward. When a user types or modifies text, the editor updates its content accordingly. However, multi-user text editing introduces complexities as it involves real-time collaboration among multiple participants.

So, what are the challenges associated with multi-user text editing?

Conflict resolution: One of the main challenges is managing conflicts that may arise when multiple users attempt to edit the same portion of the text simultaneously. For example, if two users modify the text at position X, it becomes crucial to determine whose changes should take precedence to ensure that the final result is consistent.
Real-time coordination: Ensuring a seamless user experience is another challenge. In a collaborative setting, users expect their edits to be applied immediately, similar to the responsiveness of single-user text editing. Achieving this real-time experience requires efficient handling and processing of each operation to provide smooth interactions for all participants.
Synchronization among multiple users: Maintaining synchronization of the shared content across multiple users is a significant challenge. When a new user joins the editing session, they need to receive the current state of the text, including all the modifications made by other users. Likewise, any changes made by one user must be propagated and updated for all other users in real time.

Collaborative text editors overcome the challenges of multi-user text editing by employing a range of techniques to enable real-time collaboration and synchronization.

To understand how they work, let’s try to model a basic collaborative text editor. But before we do that, let’s begin by using a simple text editor and familiarizing ourselves with the challenges mentioned earlier.

A basic "collaborative" editor.

In the above diagram, we have set up a server that sends operations to each client. However, the clients don’t implement any conflict resolution mechanism. This is the reason why the documents didn’t converge.

Divergent state is a big no-no when it comes to collaborative editing. But let’s dissect what happened:

User A adds the character “T” at position 0. This operation is done locally first, and then sent to the other client through the server. The local state is changed to “TWIN”.
Simultaneously, user B deletes the character “W” at position 0. This operation is also done locally first, and then sent to the other client through the server. The local state is changed to “IN”.
While user A had its local state as “TWIN”, it receives the operation from user B via the server. This operation is applied locally, and now the state is changed to “WIN”.
Similarly, user A’s operation is received by user B, and the local state changes from “IN” to “TIN”.

So…. how do we fix this?

Let’s add a black box to this whole situation.

Adding a black box to the scenario

This black box should handle conflict resolution and provide the ability to perform real-time collaborative editing.

The two options for this black box are:

OT (Operational transformation)
CRDTs (Conflict-free replicated datatypes)

Let’s understand each of them in detail.

Transforming operations ¶

Operational Transformation (OT) is a fundamental technique employed in collaborative editing systems to ensure consistency and concurrency control.

The basic unit of OT is an operation, and the two possible operations are: insert and delete. Any operation being performed on the document results in a new document state.

OT solves the issue of concurrent editing by transforming each operation using the previous operation, before it is applied to the document.

To understand this in a better way, let’s use OT as the black box in our previous example:

Let’s assume that we have two users A and B, where A tries to change the version number to 1.21. This change is applied locally:

User A performs some operations

Now, user B tries to add “Go” to the sentence, as the discussion between A and B is related to the Go 1.21 release; and these changes are applied locally:

User B performs some operations

Now, the first operation which A made is received by B, and is applied on its local state:

User B applies A's operations

However, as shown above, this change is not intended. A just wanted to change the Go version number, that’s it. This is why we need a synchronization algorithm - it would help A and B to communicate effectively while preserving intent.

Let’s see what happens if we use something like OT (Operational Transformation). We’ll emulate the same operation:

User B applies A's delete using OT

See? OT tries to preserve intent by transforming operations - in this case, delete(3) is transformed to delete(5) based on B’s local state. This ensures that the change is applied at the correct place. Similarly, A’s insert operation is transformed:

User B applies A's insert using OT

And this is it - A and B can now communicate effectively!

There are multiple consistency models for OT algorithms, but I want to keep this short - I’ll focus on the CCI model. The CCI model revolves around these three properties:

Convergence: Eventually all documents in the system will be in the same state. Documents diverge when a client makes an edit to a document. For an OT algorithm to be correct, the algorithm must ensure that all documents in the system eventually converge on the same state.
Causality preservation: For any two operations, namely o1 and o2, where o1 → o2, the order of execution of both operations across all clients must strictly follow the order o1 before o2.
Intention Preservation: The impact of any operation o must consistently reflect the intended effect as envisioned by the operation’s creator across all clients.

There are a bunch of other concepts (for example, forward and backward transformation) which aren’t covered here, but I would highly advise to refer the Wikipedia page for Operational Transformation if you’re more curious.

OTs are cool, but what about CRDTs? Are they any better?

Well.. let’s find out!

CRDTs ¶

A Conflict-free Replicated Data Type (CRDT) is an abstract data structure with a well-defined interface. Its purpose is to be duplicated across multiple nodes, demonstrating the following key characteristics:

Modifications can be made to any replica independently, without the need for coordination with other replicas.
When two replicas have received the same set of updates, they ensure convergence to an identical state through well-defined mathematical principles, ensuring consistent results.

CRDTs are crafted for replication across distributed systems, prioritizing eventual consistency.

To summarize, CRDTs enable updates to be applied without resource-intensive synchronization or consensus mechanisms.

And the best part? They are eventually consistent. Pretty neat, huh?

CRDTs can be categorized into two general types based on the model of replication: state-based and operation-based.

State-based CRDTs ¶

State-based CRDTs, also known as convergent replicated data types (CvRDTs), work by sending its full state over the wire. On receiving the state from a different replica, the local node merges the received state with the local state.

The main issue with state-based CRDTs is that its performance depends on the size of the data (state). If the data transmitted over the wire is big, state-based CRDTs may be inefficient.

Operation-based CRDTs ¶

Operation-based CRDTs, also known as commutative replicated data types (CmRDTs), work by sending individual operations over the wire.

In order for all replicas to converge to the same state, all the operations are required to be commutative when applied to the local state.

The main benefit of operation-based CRDTs is the minimal amount of data sent over the network as compared to state-based CRDTs.

OT vs CRDTs ¶

So, you might be wondering - what’s the real difference?

To summarize: Operational transformation (OT) works by changing operations. CRDTs work by changing state.

OT (Operational Transformation)	CRDTs (Conflict-free Replicated Data Types)
Relies on conflict resolution mechanisms	Handles conflicts by design
Requires a centralized server or coordination for conflict resolution	Conflict resolution is automatic and decentralized
Can be complex to implement	Simpler to implement
Convergence time can be longer	Achieves rapid convergence

Phew! I’ve covered quite a bit of theory, and I’d like to avoid making this longer and potentially tedious for you. Let’s dive into the actual implementation of the editor!

Building pairpad

To simplify the process of building a collaborative editor, I devised a plan to start with a basic architecture and gradually add more features while keeping the core elements straightforward.

For pairpad, the MVP needed to have:

a good interface for interactions
a smooth collaborative editing experience
the easiest deployment experience
an instant way to connect to a session

Having established these primary requirements, the next step was identifying the core components of the collaborative editor. For a collaborative editor, the components would be:

A client (for interaction)
An event handler (which would respond to user input and perform actions)
A conflict resolution mechanism (OT/CRDTs)
A broadcaster (which would broadcast operations to other clients)
anything else..

By using these components as building blocks, I came up with this architecture for pairpad:

High-level architecture for pairpad

I’ll also attach the original architecture diagram:

Original architecture diagram (Oct 2022)

The diagram doesn’t explain much; let me explain each flow in depth:

User interface ¶

The user interface is a text-based user interface (TUI) built using the Go port of termbox, a minimal alternative to ncurses.

Each keystroke is consumed by termbox and is registered as a termbox event. This event is sent to a channel called termboxChan.

This channel is the main “pipeline” of any UI-related events. All events from this channel are consumed by the “key event handler”, which processes the event and performs operations based on the key. Operations include cursor movements, performing editor operations and payload generation, which includes the generation of a JSON payload, which is sent to the client’s connection engine.

UI architecture diagram.

Client ¶

The payload sent from the event handler is now parsed by the connection engine, the component which is responsible for communicating with the server. A goroutine keeps listening for any messages from the WebSocket connection, and sends the message to the message channel (messageChan).

These messages are parsed by the message handler, which checks the message type and performs a suitable action on the local document (CRDT). This will be discussed in detail later.

Client architecture diagram.

Server ¶

Finally, the server is just a hub that broadcasts messages to all connected clients. All clients connect to the server using WebSockets, and the server maintains a list of all connected clients.

The server broadcasts all operations sent by a client to all connected clients. It also helps in managing a document synchronization request, wherein a client asks for the current copy of the editing session in order to be in sync with other clients.

Server architecture.

Let me also address some potential questions that may arise in your mind:

Q: Why did you choose CRDTs over OT?

Pretty simple - I always wanted to tinker around with CRDTs, and OT seemed pretty complicated at the first glance. Since this was a toy project and not a commercial project, I decided to go with the approach that seemed simpler to me, and CRDTs were magical to me - the idea of handling conflicts by design rather than resolution seemed fascinating to me.

Q: Why did you choose a client-server architecture over a P2P architecture?

A: I wanted the user flow to be really simple: all clients should be able to connect to a central server, and users can maintain their own server if needed. I wanted to keep the amount of maintainance and work needed to set up a session to be minimal. And for keeping it more accessible for users who are less experienced with server maintainance, this decision was made. I’m not saying that setting up P2P apps is hard, but when you develop software, you need to ensure that it’s accessible for everyone.

Setting up the user interface ¶

It was time to pick a library for building the TUI (user interface). I had three options:

bubbletea and the charm.sh ecosystem
termbox
tcell

The charm.sh ecosystem is a no-brainer if you want to build beautiful applications for the terminal.

I played around with bubbletea, but it hit me that whipping up the TUI I wanted might be trickier than I thought. The fine-grained control (over cursor positions, individual cells, etc.) was not quite there. And even if there was a workaround, keeping up with how to pull it off started feeling like a bit of a struggle, to be honest.

I did ask the team about this, and here was the response from the awesome @muesli:

Screenshot of a conversation in the charm.sh Discord

I tried my best, but there was definitely something I was lacking during the implementation – it felt like quite the challenge.

If any bubbletea enthusiast could help in porting the current TUI from termbox to bubbletea through some tips or code contributions, I’d be happy to listen!

As I didn’t want this to be a blocker, so I scrapped the prototype and moved on to the next option: termbox - a classic solution.

termbox uses a two-dimensional cell buffer to represent the terminal screen, with a few methods like setting cells, clearing the screen, or setting the cursor. All input (including keyboard, mouse, or resize events) is tracked through a unified Event type.

This is exactly what I needed, and it makes a great fit for applications like text editors. The biggest advantage of using termbox was the ability to utilize methods for low-level functionality too. (setting cells/cursors, etc.)

The editor UI is structured in this way:

Editor UI representation

The UI is pretty simplistic, with only two components:

Core Editing View
Status bar

The core editing view is where the actual text editing happens. Users can type, delete, move the cursor, select text, etc. in this view. How this works is by maintaining a buffer of cells representing the terminal screen. Each cell contains either a character or a style (like bold, underline, etc.). On modification, the buffer is updated and redrawn on the screen.

The status bar is a component made using the termbox buffer which acts as a place for displaying information like cursor position, info/error messages, etc. It is updated whenever there are relevant changes in the core editing view, or when there are events like user joining/leaving a session, etc.

And here is the editor implementation using termbox-go (pure Go implementation of termbox):

package editor

import (
	"fmt"
	"sync"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
)

type EditorConfig struct {
	ScrollEnabled bool
}

// Editor represents the editor's skeleton.
// The editor is composed of two components:
// 1. an editable text area; which acts as the primary interactive area.
// 2. a status bar; which displays messages on different events, for example, when an user joins a session, etc.
type Editor struct {
	// Text contains the editor's content.
	Text []rune

	// Cursor represents the cursor position of the editor.
	Cursor int

	// Width represents the terminal's width in characters.
	Width int

	// Height represents the terminal's width in characters.
	Height int

	// ColOff is the number of columns between the start of a line and the left of the editor window
	ColOff int

	// RowOff is the number of rows between the beginning of the text and the top of the editor window
	RowOff int

	// ShowMsg acts like a switch for the status bar.
	ShowMsg bool

	// StatusMsg holds the text to be displayed in the status bar.
	StatusMsg string

	// StatusChan is used to send and receive status messages.
	StatusChan chan string

	// StatusMu protects against concurrent reads and writes to status bar info.
	StatusMu sync.Mutex

	// Users holds the names of all users connected to the server, displayed in the status bar.
	Users []string

	// ScrollEnabled determines whether or not the user can scroll past the initial editor
	// window. It is set by the EditorConfig.
	ScrollEnabled bool

	// IsConnected shows whether the editor is currently connected to the server.
	IsConnected bool

	// DrawChan is used to send and receive signals to update the terminal display.
	DrawChan chan int

	// mu prevents concurrent reads and writes to the editor state.
	mu sync.RWMutex
}

var userColors = []termbox.Attribute{
	termbox.ColorGreen,
	termbox.ColorYellow,
	termbox.ColorBlue,
	termbox.ColorMagenta,
	termbox.ColorCyan,
	termbox.ColorLightYellow,
	termbox.ColorLightMagenta,
	termbox.ColorLightGreen,
	termbox.ColorLightRed,
	termbox.ColorRed,
}

// NewEditor returns a new instance of the editor.
func NewEditor(conf EditorConfig) *Editor {
	return &Editor{
		ScrollEnabled: conf.ScrollEnabled,
		StatusChan:    make(chan string, 100),
		DrawChan:      make(chan int, 10000),
	}
}

// GetText returns the editor's content.
func (e *Editor) GetText() []rune {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.Text
}

// SetText sets the given string as the editor's content.
func (e *Editor) SetText(text string) {
	e.mu.Lock()
	e.Text = []rune(text)
	e.mu.Unlock()
}

// GetX returns the X-axis component of the current cursor position.
func (e *Editor) GetX() int {
	x, _ := e.calcXY(e.Cursor)
	return x
}

// SetX sets the X-axis component of the current cursor position to the specified X position.
func (e *Editor) SetX(x int) {
	e.Cursor = x
}

// GetY returns the Y-axis component of the current cursor position.
func (e *Editor) GetY() int {
	_, y := e.calcXY(e.Cursor)
	return y
}

....

// MoveCursor updates the cursor position horizontally by a given x increment, and
// vertically by one line in the direction indicated by y. The positive directions are
// right and down, respectively.
// This is used by the UI layer, where it updates the cursor position on keypresses.
func (e *Editor) MoveCursor(x, y int) {
	if len(e.Text) == 0 && e.Cursor == 0 {
		return
	}
	// Move cursor horizontally.
	newCursor := e.Cursor + x

	// Move cursor vertically.
	if y > 0 {
		newCursor = e.calcCursorDown()
	}

	if y < 0 {
		newCursor = e.calcCursorUp()
	}

	if e.ScrollEnabled {
		cx, cy := e.calcXY(newCursor)

		// move the window to adjust for the cursor
		rowStart := e.GetRowOff()
		rowEnd := e.GetRowOff() + e.GetHeight() - 1

		if cy <= rowStart { // scroll up
			e.IncRowOff(cy - rowStart - 1)
		}

		if cy > rowEnd { // scroll down
			e.IncRowOff(cy - rowEnd)
		}

		colStart := e.GetColOff()
		colEnd := e.GetColOff() + e.GetWidth()

		if cx <= colStart { // scroll left
			e.IncColOff(cx - (colStart + 1))
		}

		if cx > colEnd { // scroll right
			e.IncColOff(cx - colEnd)
		}
	}

	// Reset to bounds.
	if newCursor > len(e.Text) {
		newCursor = len(e.Text)
	}

	if newCursor < 0 {
		newCursor = 0
	}
	e.mu.Lock()
	e.Cursor = newCursor
	e.mu.Unlock()
}

Implementing the CRDT ¶

With the UI in place, it was time to find a suitable CRDT. Following a thorough exploration of the CRDT landscape, I opted for WOOT. WOOT is simple enough and works well for collaborative text editing. Other potential options included Logoot, etc. but WOOT seemed the best fit considering that I wanted to get a basic implementation up and running quickly.

WOOT ¶

WOOT is a simple sequence CRDT. In WOOT, each character is represented as a W-character which is a tuple consisting of five fields:

< id, visible, value, idcp, idcn >

where:

id is the identifier of the character (site id, unique natural number).
visible, a boolean value that indicates if the character is visible.
value is the alphabetical value of the character.
idcp is the identifier of the previous W-character of c.
idcn is the identifier of the next W-character of c.

Each site has an unique identifier s, a logical clock, and a sequence of W-characters. The site identifier and the local clock are used for constructing the unique identifier for the characters.

Each time a W-character is generated on site s, the local clock is incremented. This is how the uniqueness of the character identifiers is ensured.

Local insertions are pretty straightforward - a new W-character is created with idcp and idcn pointing towards the previous and the next character.

Integrating insertions from other sites is a bit complicated - WOOT will take a subsequence of the characters between idcp and idcn, if the subsequence is empty, it can be inserted here. Else, it will try to find a suitable position and make a recursive call to itself with those positions.

And for deletions, instead of directly removing characters, WOOT marks the character as deleted by setting the visible boolean to false.

To start off, let’s implement an interface for abstracting away the underlying CRDT. This would ensure that we can keep the CRDT layer separate from the CRDT itself.

type CRDT interface {
	Insert(position int, value string) (string, error)
	Delete(position int) string
}

For implementing WOOT, I started porting psuedocode from the paper to actual Go code. I took some references from various JavaScript/TypeScript implementations while writing this. If you want to read more about WOOT, you can read the research paper here.

Here is the complete implementation of WOOT:

// LocalInsert inserts the character into the document.
func (doc *Document) LocalInsert(char Character, position int) (*Document, error) {
	if position <= 0 || position >= doc.Length() {
		return doc, ErrPositionOutOfBounds
	}

	if char.ID == "" {
		return doc, ErrEmptyWCharacter
	}

	doc.Characters = append(doc.Characters[:position],
		append([]Character{char}, doc.Characters[position:]...)...,
	)

	// Update next and previous pointers.
	doc.Characters[position-1].IDNext = char.ID
	doc.Characters[position+1].IDPrevious = char.ID

	return doc, nil
}

// IntegrateInsert inserts the given Character into the Document
// Characters based off of the previous & next Character
func (doc *Document) IntegrateInsert(char, charPrev, charNext Character) (*Document, error) {
	// Get the subsequence.
	subsequence, err := doc.Subseq(charPrev, charNext)
	if err != nil {
		return doc, err
	}

	// Get the position of the next character.
	position := doc.Position(charNext.ID)
	position--

	// If no characters are present in the subsequence, insert at current position.
	if len(subsequence) == 0 {
		return doc.LocalInsert(char, position)
	}

	// If one character is present in the subsequence, insert at previous position.
	if len(subsequence) == 1 {
		return doc.LocalInsert(char, position-1)
	}

	// Find the appropriate position for the new character in the subsequence.
    i := 1
    for i < len(subsequence)-1 && subsequence[i].ID < char.ID {
        i++
    }

    // Make a recursive call with the correct positions.
    return doc.IntegrateInsert(char, subsequence[i-1], subsequence[i])
}

// GenerateInsert generates a character for a given value.
func (doc *Document) GenerateInsert(position int, value string) (*Document, error) {
	// Increment local clock.
	mu.Lock()
	LocalClock++
	mu.Unlock()

	// Get previous and next characters.
	charPrev := IthVisible(*doc, position-1)
	charNext := IthVisible(*doc, position)

	// Use defaults.
	if charPrev.ID == "-1" {
		charPrev = doc.Find("start")
	}
	if charNext.ID == "-1" {
		charNext = doc.Find("end")
	}

	char := Character{
		ID:         fmt.Sprint(SiteID) + fmt.Sprint(LocalClock),
		Visible:    true,
		Value:      value,
		IDPrevious: charPrev.ID,
		IDNext:     charNext.ID,
	}

	return doc.IntegrateInsert(char, charPrev, charNext)
}

// IntegrateDelete finds a character and marks it for deletion.
func (doc *Document) IntegrateDelete(char Character) *Document {
	position := doc.Position(char.ID)
	if position == -1 {
		return doc
	}

	// This is how deletion is done.
	doc.Characters[position-1].Visible = false

	return doc
}

// GenerateDelete generates the character which is to be marked for deletion.
func (doc *Document) GenerateDelete(position int) *Document {
	char := IthVisible(*doc, position)
	return doc.IntegrateDelete(char)
}

Now, WOOT needs to implement the CRDT interface - pretty straightforward:

func (doc *Document) Insert(position int, value string) (string, error) {
	newDoc, err := doc.GenerateInsert(position, value)
	if err != nil {
		return Content(*doc), err
	}

	return Content(*newDoc), nil
}

func (doc *Document) Delete(position int) string {
	newDoc := doc.GenerateDelete(position)
	return Content(*newDoc)
}

Building the client ¶

With the editor implementation in place, the next step is to set up the TUI and the event handler.

The event handler is simple: for each termbox event, we need to perform some action. In our case, we have multiple event-action mappings:

On Esc, the editor should exit
On Ctrl+S, the document should be persisted to a local file
On Ctrl+L, the document should load the file
etc.

Here is the code for the termbox event handler:

func handleTermboxEvent(ev termbox.Event, conn *websocket.Conn) error {
	// We only want to deal with termbox key events (EventKey).
	if ev.Type == termbox.EventKey {
		switch ev.Key {

		// The default keys for exiting an session are Esc and Ctrl+C.
		case termbox.KeyEsc, termbox.KeyCtrlC:
			// Return an error with the prefix "pairpad", so that it gets treated as an exit "event".
			return errors.New("pairpad: exiting")

		// The default key for saving the editor's contents is Ctrl+S.
		case termbox.KeyCtrlS:
			// If no file name is specified, set filename to "pairpad-content.txt"
			if fileName == "" {
				fileName = "pairpad-content.txt"
			}

			// Save the CRDT to a file.
			err := crdt.Save(fileName, &doc)
			if err != nil {
				logrus.Errorf("Failed to save to %s", fileName)
				e.StatusChan <- fmt.Sprintf("Failed to save to %s", fileName)
				return err
			}

			// Set the status bar.
			e.StatusChan <- fmt.Sprintf("Saved document to %s", fileName)

		// The default key for loading content from a file is Ctrl+L.
		case termbox.KeyCtrlL:
			if fileName != "" {
				logger.Log(logrus.InfoLevel, "LOADING DOCUMENT")
				newDoc, err := crdt.Load(fileName)
				if err != nil {
					logrus.Errorf("failed to load file %s", fileName)
					e.StatusChan <- fmt.Sprintf("Failed to load %s", fileName)
					return err
				}
				e.StatusChan <- fmt.Sprintf("Loading %s", fileName)
				doc = newDoc
				e.SetX(0)
				e.SetText(crdt.Content(doc))

				logger.Log(logrus.InfoLevel, "SENDING DOCUMENT")
				docMsg := commons.Message{Type: commons.DocSyncMessage, Document: doc}
				_ = conn.WriteJSON(&docMsg)
			} else {
				e.StatusChan <- "No file to load!"
			}

		// The default keys for moving left inside the text area are the left arrow key, and Ctrl+B (move backward).
		case termbox.KeyArrowLeft, termbox.KeyCtrlB:
			e.MoveCursor(-1, 0)

		// The default keys for moving right inside the text area are the right arrow key, and Ctrl+F (move forward).
		case termbox.KeyArrowRight, termbox.KeyCtrlF:
			e.MoveCursor(1, 0)

		// The default keys for moving up inside the text area are the up arrow key, and Ctrl+P (move to previous line).
		case termbox.KeyArrowUp, termbox.KeyCtrlP:
			e.MoveCursor(0, -1)

		// The default keys for moving down inside the text area are the down arrow key, and Ctrl+N (move to next line).
		case termbox.KeyArrowDown, termbox.KeyCtrlN:
			e.MoveCursor(0, 1)

		// Home key, moves cursor to initial position (X=0).
		case termbox.KeyHome:
			e.SetX(0)

		// End key, moves cursor to final position (X= length of text).
		case termbox.KeyEnd:
			e.SetX(len(e.Text))

		// The default keys for deleting a character are Backspace and Delete.
		case termbox.KeyBackspace, termbox.KeyBackspace2:
			performOperation(OperationDelete, ev, conn)
		case termbox.KeyDelete:
			performOperation(OperationDelete, ev, conn)

		// The Tab key inserts 4 spaces to simulate a "tab".
		case termbox.KeyTab:
			for i := 0; i < 4; i++ {
				ev.Ch = ' '
				performOperation(OperationInsert, ev, conn)
			}

		// The Enter key inserts a newline character to the editor's content.
		case termbox.KeyEnter:
			ev.Ch = '\n'
			performOperation(OperationInsert, ev, conn)

		// The Space key inserts a space character to the editor's content.
		case termbox.KeySpace:
			ev.Ch = ' '
			performOperation(OperationInsert, ev, conn)

		// Every other key is eligible to be a candidate for insertion.
		default:
			if ev.Ch != 0 {
				performOperation(OperationInsert, ev, conn)
			}
		}
	}

	e.SendDraw()
	return nil
}

performOperation is used for performing operations on the CRDT (local document), and it sends a payload over the WebSocket connection.

const (
	OperationInsert = iota
	OperationDelete
)

// performOperation performs a CRDT insert or delete operation on the local document and sends a message over the WebSocket connection.
func performOperation(opType int, ev termbox.Event, conn *websocket.Conn) {
	// Get position and value.
	ch := string(ev.Ch)

	var msg commons.Message

	// Modify local state (CRDT) first.
	switch opType {
	case OperationInsert:
		logger.Infof("LOCAL INSERT: %s at cursor position %v\n", ch, e.Cursor)

		text, err := doc.Insert(e.Cursor+1, ch)
		if err != nil {
			e.SetText(text)
			logger.Errorf("CRDT error: %v\n", err)
		}
		e.SetText(text)

		e.MoveCursor(1, 0)
		msg = commons.Message{Type: "operation", Operation: commons.Operation{Type: "insert", Position: e.Cursor, Value: ch}}

	case OperationDelete:
		logger.Infof("LOCAL DELETE: cursor position %v\n", e.Cursor)

		if e.Cursor-1 < 0 {
			e.Cursor = 0
		}

		text := doc.Delete(e.Cursor)
		e.SetText(text)

		msg = commons.Message{Type: "operation", Operation: commons.Operation{Type: "delete", Position: e.Cursor}}
		e.MoveCursor(-1, 0)
	}

	// Send the message.
	if e.IsConnected {
		err := conn.WriteJSON(msg)
		if err != nil {
			e.IsConnected = false
			e.StatusChan <- "lost connection!"
		}
	}
}

How do the other clients handle the received payload? Here’s an insight into the client’s internal process:

handleMsg receives the message from a top-level component (I’ll talk about this later)
it checks the message type; pairpad has multiple message types: one for operations, one for info messages, etc.
it performs an action based on the message type

You might be seeing DocSync, DocReq and a bunch of other types here. Those are for document synchronization, I’ll discuss that later.

Here is the code for handling a message:

// handleMsg updates the CRDT document with the contents of the message.
func handleMsg(msg commons.Message, conn *websocket.Conn) {
	switch msg.Type {
	case commons.DocSyncMessage:
		logger.Infof("DOCSYNC RECEIVED, updating local doc %+v\n", msg.Document)

		doc = msg.Document
		e.SetText(crdt.Content(doc))

	case commons.DocReqMessage:
		logger.Infof("DOCREQ RECEIVED, sending local document to %v\n", msg.ID)

		docMsg := commons.Message{Type: commons.DocSyncMessage, Document: doc, ID: msg.ID}
		_ = conn.WriteJSON(&docMsg)

	case commons.SiteIDMessage:
		siteID, err := strconv.Atoi(msg.Text)
		if err != nil {
			logger.Errorf("failed to set siteID, err: %v\n", err)
		}

		crdt.SiteID = siteID
		logger.Infof("SITE ID %v, INTENDED SITE ID: %v", crdt.SiteID, siteID)

	case commons.JoinMessage:
		e.StatusChan <- fmt.Sprintf("%s has joined the session!", msg.Username)

	case commons.UsersMessage:
		e.StatusMu.Lock()
		e.Users = strings.Split(msg.Text, ",")
		e.StatusMu.Unlock()

	default:
		switch msg.Operation.Type {
		case "insert":
			_, err := doc.Insert(msg.Operation.Position, msg.Operation.Value)
			if err != nil {
				logger.Errorf("failed to insert, err: %v\n", err)
			}

			e.SetText(crdt.Content(doc))
			if msg.Operation.Position-1 <= e.Cursor {
				e.MoveCursor(len(msg.Operation.Value), 0)
			}
			logger.Infof("REMOTE INSERT: %s at position %v\n", msg.Operation.Value, msg.Operation.Position)

		case "delete":
			_ = doc.Delete(msg.Operation.Position)
			e.SetText(crdt.Content(doc))
			if msg.Operation.Position-1 <= e.Cursor {
				e.MoveCursor(-len(msg.Operation.Value), 0)
			}
			logger.Infof("REMOTE DELETE: position %v\n", msg.Operation.Position)
		}
	}

	// printDoc is used for debugging purposes. Don't comment this out.
	// This can be toggled via the `-debug` flag.
	// The default behavior for printDoc is to NOT log anything.
	// This is to ensure that the debug logs don't take up much space on the user's filesystem, and can be toggled on demand.
	printDoc(doc)

	e.SendDraw()
}

You may have observed that communication takes place across several channels. Here’s a breakdown of the functions of these channels:

termboxChan is a channel meant for termbox events. A goroutine sends termbox events to the channel by using the PollEvent method, which polls for events and returns a termbox event (termbox.Event).

// getTermboxChan returns a channel of termbox Events repeatedly waiting on user input.
func getTermboxChan() chan termbox.Event {
	termboxChan := make(chan termbox.Event)

	go func() {
		for {
			termboxChan <- termbox.PollEvent()
		}
	}()

	return termboxChan
}

messageChan is a channel meant for internal “messages”. A goroutine reads JSON from the connection, and sends the message to the message channel.

// getMsgChan returns a message channel that repeatedly reads from a websocket connection.
func getMsgChan(conn *websocket.Conn) chan commons.Message {
	messageChan := make(chan commons.Message)
	go func() {
		for {
			var msg commons.Message

			// Read message.
			err := conn.ReadJSON(&msg)
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					logger.Errorf("websocket error: %v", err)
				}
				e.IsConnected = false
				e.StatusChan <- "lost connection!"
				break
			}

			logger.Infof("message received: %+v\n", msg)

			// send message through channel
			messageChan <- msg

		}
	}()
	return messageChan
}

With of all these communication mechanisms in place, it’s time to look at the main UI loop:

initUI starts termbox and creates a new editor view
initUI runs two goroutines:
- handleStatusMsg for updating the status bar
- drawLoop for updating the views
initUI calls the main loop
mainLoop has a never-ending for loop, which focuses on two things:
- listen for a termbox Event from termboxChan, and pass the event to handleTermboxEvent
- listen for a message from msgChan and pass the message to handleMsg

// initUI creates a new editor view and runs the main loop.
func initUI(conn *websocket.Conn, conf UIConfig) error {
	err := termbox.Init()
	if err != nil {
		return err
	}
	defer termbox.Close()

	e = editor.NewEditor(conf.EditorConfig)
	e.SetSize(termbox.Size())
	e.SetText(crdt.Content(doc))
	e.SendDraw()
	e.IsConnected = true

	go handleStatusMsg()

	go drawLoop()

	err = mainLoop(conn)
	if err != nil {
		return err
	}

	return nil
}

// mainLoop is the main update loop for the UI.
func mainLoop(conn *websocket.Conn) error {
	// termboxChan is used for sending and receiving termbox events.
	termboxChan := getTermboxChan()

	// msgChan is used for sending and receiving messages.
	msgChan := getMsgChan(conn)

	for {
		select {
		case termboxEvent := <-termboxChan:
			err := handleTermboxEvent(termboxEvent, conn)
			if err != nil {
				return err
			}
		case msg := <-msgChan:
			handleMsg(msg, conn)
		}
	}
}

Building the server ¶

pairpad employs a simple mechanism for broadcasting. The idea is to use the server as a hub that broadcasts messages to all connected clients. All clients connect to the server using WebSockets, and the server maintains a list of all connected clients.

On each operation, the client sends a message to the server with the operation details. The server then broadcasts that message to all connected clients. Each client applies that operation to their local document, keeping all documents in sync.

Since sessions are long-lived, using WebSockets was a no-brainer over plain TCP or HTTP polling. WebSockets allow for bidirectional, full-duplex communication between clients and servers. This means that both the server and clients can push messages to each other at any time.

When a client connects to the server, we generate a site ID for that client and add that client to the list of connected clients. We send a site ID “message” to that client to let it know its ID. This site ID is used by the client for constructing the unique identifier for a character: <SiteID><LocalClock>. In the diagram, character ID is represented as <SiteID>.<LocalClock> for better readability.

Site ID generation

func handleConn(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		color.Red("Error upgrading connection to websocket: %v\n", err)
		conn.Close()
		return
	}
	defer conn.Close()

	clientID := uuid.New()

	// Carefully increment and assign site ID with mutexes.
	mu.Lock()
	siteID++

	client := &client{
		Conn:    conn,
		SiteID:  strconv.Itoa(siteID),
		id:      clientID,
		writeMu: sync.Mutex{},
		mu:      sync.Mutex{},
	}
	mu.Unlock()

	clients.add(client)

	siteIDMsg := commons.Message{Type: commons.SiteIDMessage, Text: client.SiteID, ID: clientID}
	clients.broadcastOne(siteIDMsg, clientID)

	docReq := commons.Message{Type: commons.DocReqMessage, ID: clientID}
	clients.broadcastOneExcept(docReq, clientID)

	clients.sendUsernames()

	// Read messages from the connection and send to channel to broadcast
	for {
		var msg commons.Message
		if err := client.read(&msg); err != nil {
			color.Red("Failed to read message. closing client connection with %s. Error: %s", client.Username, err)
			return
		}

		// Send docSync to handleSync function. DocSync message IDs refer to
		// their destination. This channel send should happen before reassigning the
		// msg.ID
		if msg.Type == commons.DocSyncMessage {
			syncChan <- msg
			continue
		}

		// Set message ID as the ID of the sending client. Most message IDs refer to
		// their origin.
		msg.ID = clientID

		// Send message to messageChan for logging and broadcasting
		messageChan <- msg
	}
}

// handleMsg listens to the messageChan channel and broadcasts messages to other clients.
func handleMsg() {
	for {
		// Get message from messageChan.
		msg := <-messageChan

		// Log each message to stdout.
		t := time.Now().Format(time.ANSIC)
		if msg.Type == commons.JoinMessage {
			clients.updateName(msg.ID, msg.Username)
			color.Green("%s >> %s %s (ID: %s)\n", t, msg.Username, msg.Text, msg.ID)
			clients.sendUsernames()
		} else if msg.Type == "operation" {
			color.Green("operation >> %+v from ID=%s\n", msg.Operation, msg.ID)
		} else {
			color.Green("%s >> unknown message type:  %v\n", t, msg)
			clients.sendUsernames()
			continue
		}

		clients.broadcastAllExcept(msg, msg.ID)
	}
}

// handleSync reads from the syncChan and sends the message to the appropriate user(s).
func handleSync() {
	for {
		syncMsg := <-syncChan
		switch syncMsg.Type {
		case commons.DocSyncMessage:
			clients.broadcastOne(syncMsg, syncMsg.ID)
		case commons.UsersMessage:
			color.Blue("usernames: %s", syncMsg.Text)
			clients.broadcastAll(syncMsg)
		}
	}
}

Implementing a document sync “protocol” ¶

Remember seeing things like docReq and docSync in the above code snippets? Those were custom message types used to synchronize documents between clients.

When a client connects to the server, it has no idea on the existing state of the other clients, so it sends a docReq message to request the current document state from one of the other clients.

Once the docReq message is received by the client, it responds by sending a docSync message containing the entire document content. This message is sent to the client that requested the document, helping it sync to the current state.

This protocol might have some flaws, but it works well in most cases, and serves as a simple approach to synchronizing documents between clients.

Here is a simple overview of the protocol, explained through diagrams:

Client 2 disconnects

Client 3 tries sync

Client 3 sends docReq

Server checks for active clients and forwards docReq

Client 1 receives docReq, and sends docSync

Client 3 receives content and merges it locally

Testing it out! ¶

For testing pairpad, I packaged it in a Docker image and used fly.io for deployments. Fly is pretty simple to get started with, and the free tier is more than enough for testing purposes.

Connecting to the server is pretty straightforward, all you need to do is:

pairpad -server pairpad.test

where pairpad.test is your server location.

I’ve tested pairpad multiple times locally, but testing it on a remote server was pretty cool! It felt like I actually built something useful that could be used by others. I was expecting latency issues and synchronization problems, but the test session went pretty well!

Testing pairpad from two different continents!

(I don’t have a recording of the actual session, but if you’re curious, I highly recommend you to try it out yourself!)

Is it perfect? ¶

It definitely isn’t. pairpad started out as a fun project, and still isn’t close to completion. Software grows eternally, and there’s always room for improvement. You might even find multiple occurrences of inefficient code.

Having said that, the project is open for contributions! All kinds of contributions are welcome - code, documentation, general suggestions, etc.

Here’s a list of things which I wish pairpad could have:

A web client
Ability to create private rooms
Telescope-like navigation system
Workspaces (handle multiple files)

If you’d like to work on any of these features (or more), please consider raising a pull request!

Acknowledgements ¶

Special thanks to Ben for making this possible - I had a great time pairing with Ben on this one! I would also like to thank these awesome people:

Miccah, for helping out with code structure
Nicole, for nerd-sniping and finding a bug that was messing up synchronization
Henrique, for providing insights for building a web-based version of the editor.

Reverse image search using perceptual hashes

2022-12-10T00:00:00Z

Using perceptual hashes for image search.

I was curious about how Google Image Search worked, but I didn’t want to delve right into their implementation. Instead, I wanted to explore how to build my own reverse image search engine from scratch.

I had a rough sketch ready: user puts in their image, the engine searches for an image from a “database” and returns relevant items back to the user. But, what is the engine, really?

For checking if an item is similar to another, you need to have something to compare against it. Images have pixels, but is it the best choice? Nope. It’s really obvious that it would only work well on completely identical images, and not similar images. The emphasis is on similarity.

One option is to represent an image as something more tangible, such as a hash. However, using cryptographic hashes like MD5 or SHA256 is not suitable for this context, as they are designed to generate unpredictable results. Small changes to an image could cause drastic changes in the hash, making it unsuitable for comparing similar images.

That’s when I discovered perceptual hashes, which felt like a promising starting point.

But what even are perceptual hashes?

Perceptual hashes ¶

Perceptual hashes are a type of digital fingerprinting algorithm that produces a hash, or a fingerprint that uniquely identifies the input media, just like human fingerprints. Unlike traditional cryptographic hashes, which are designed to produce a unique and unpredictable output for any given input, perceptual hashes are specifically tailored to capture the unique visual characteristics of an image in a way that is robust to minor variations.

Perceptual hashing algorithms work by analyzing the structural and spatial characteristics of an image, such as its edges, lines, textures, and colors, and then transforming this information into a compact and standardized digital signature that can be used to identify similar images. The resulting perceptual hash is a fixed-length binary code that is highly sensitive to changes in the visual content of an image, but relatively insensitive to changes in image format, resolution, compression, or noise.

Perceptual hashes have a wide range of applications, including content-based image retrieval, image deduplication, copyright infringement detection, and reverse image search. By using perceptual hashes, it’s possible to efficiently compare large volumes of images and identify similar or identical images, even if they have been modified, resized, cropped, or otherwise altered.

Starting with ahash ¶

One of the most common perceptual hashing algorithms used for reverse image search is the average hash, or ahash for short. Here’s how it works:

Resize the input image to a fixed resolution, such as 8x8 pixels.
Convert the resized image to grayscale.
Compute the average brightness of the grayscale image.
Generate a hash by comparing each pixel’s brightness to the average value.
If the pixel is brighter than the average, mark it as a 1; if it is darker, mark it as a 0.
The resulting hash is a 64-bit binary code that can be used to compare the input image to other images in a database.

Input image

Resized image (original: 8x8, the above image is for representation (256x256))

Resized, grayscale image (original: 8x8, the above image is for representation (256x256))

By converting the input image to grayscale and resizing it to a fixed resolution, the ahash algorithm standardizes the visual content of the image and reduces the impact of irrelevant features like color or texture. By computing the average brightness of the resized image and comparing each pixel to this value, the ahash algorithm generates a binary code that is highly sensitive to changes in the visual content of the image, but relatively insensitive to changes in format, resolution, compression, or noise.

Here’s the Go code for ahash:

package hash

import (
	"image"
)

type AvgHash struct{}

func NewAverageHash() *AvgHash {
	return &AvgHash{}
}

// GenerateHash generates the hash value using average hash.
func (ah *AvgHash) GenerateHash(img image.Image) uint64 {
	// Resize source image to 8x8.
	resizedImg := Resize(8, 8, img)

	// Convert the resized image to grayscale.
	grayImg := Grayscale(resizedImg)

	// Find the mean using the grayscale image.
	mean := FindMean(grayImg)

	// Return average hash.
	return ah.getHash(grayImg, mean)
}

// getHash generates the average hash from the grayscale image and the mean.
func (*AvgHash) getHash(img image.Image, mean uint32) uint64 {
	var hash, pos uint64
	pos = 1
	bounds := img.Bounds()

	for x := bounds.Min.X; x < bounds.Max.X; x++ {
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			r, _, _, _ := img.At(x, y).RGBA()
			if r > mean {
				hash |= 1 << pos // set bit to 1
			}
			pos++
		}
	}

	return hash
}

Reverse image search ¶

Reverse image search involves finding images that are similar to a given query image.

On composing on the earlier findings, the core “engine” would look like this:

user sends an image for search
generate ahash for the image
store ahash in the database
compare the input image’s ahash to the hashes in the database

Everything is straightforward, except the comparison technique. How do you compare images based on the ahash?

Hmm, a hash is represented as a string/sequence - which string/sequence comparison technique is easier? Some sort of a technique which would give us a distance representing the difference between two entities would be a nice start.

Hmm, distance. What would be a nice algorithm for calculating distances?

Hamming distance comes to my mind - it calculates the number of bits different between the two values. Looks perfect for this usecase.

Here’s an example, the hamming distance for these hashes will be 1 (only 1 bit is different):

0111010010010010110011001000111010100010100010000100110110101010
0111010010010010110011001000011010100010100010000100110110101010

Since the comparison technique is finalized, here’s the overall flow:

user sends an image for search
generate ahash for the image
store ahash in the database
compare the input image’s ahash to the hashes in the database using hamming distance

Looks good! Here’s the Go code for the engine’s core:

(The project is at github.com/burntcarrot/hashsearch.)

package image

import (
	"github.com/burntcarrot/hashsearch/entity"
	"github.com/burntcarrot/hashsearch/pkg/hash"
)

type Service struct {
	repo Repository
}

// NewService returns the image service.
func NewService(r Repository) *Service {
	return &Service{
		repo: r,
	}
}

func (s *Service) Create(path string) error {
	// Create image entity.
	e, err := entity.NewImage(path)
	if err != nil {
		return err
	}

	// Load as image.
	img, err := hash.LoadAsImage(e.Path)
	if err != nil {
		return err
	}

	// Set the entity's hash value using the loaded image.
	ah := hash.NewAverageHash()
	e.Hash = ah.GenerateHash(img)

	return s.repo.Create(e)
}

func (s *Service) List() ([]*entity.Image, error) {
	// Get all images.
	images, err := s.repo.List()
	if err != nil {
		return nil, err
	}

	if len(images) == 0 {
		return nil, entity.ErrNotFound
	}

	return images, nil
}

func (s *Service) GetDistances(path string) ([]*entity.Image, error) {
	// Load source image.
	img, err := hash.LoadAsImage(path)
	if err != nil {
		return nil, err
	}

	// Get the source image's hash.
	ah := hash.NewAverageHash()
	sourceImageHash := ah.GenerateHash(img)

	// Get all images.
	images, err := s.repo.List()
	if err != nil {
		return nil, err
	}
	if len(images) == 0 {
		return nil, entity.ErrNotFound
	}

	// Calculate distances for each image with respect to the source image.
	for _, img := range images {
		img.Distance = hash.CalculateDistance(sourceImageHash, img.Hash)
	}

	return images, nil
}

Results ¶

Search results, sorted by distance:

The query image is the distorted image (paper-bag-cat-distorted.png):

Image with distortion and text

Here is the JSON returned from the hashsearch API:

❯ curl -s --location --request POST 'localhost:8081/v1/search' --form 'file=@"paper-bag-cat-distorted.png"' | jq
[
  {
    "path": "files/paper-bag-cat-distorted.png",
    "distance": 0,
    "hash": "0000111000011110001111100001111000000110000011100000111000111110"
  },
  {
    "path": "files/paper-bag-cat-saturated.png",
    "distance": 7,
    "hash": "0000111000001110000011100000111000101110000011100001111000111110"
  },
  {
    "path": "files/paper-bag-cat-text.png",
    "distance": 7,
    "hash": "0000111000001110000011100000111000101110000011100001111000111110"
  },
  {
    "path": "files/paper-bag-cat.png",
    "distance": 7,
    "hash": "0000111000001110000011100000111000101110000011100001111000111110"
  },
  {
    "path": "files/cat.png",
    "distance": 36,
    "hash": "0000000001111100111111001111100011111000111110000111110000000000"
  }
]

Here’s a better representation:

Image	Distance	Hash
(input image (distorted + text))	0	0000111000011110001111100001111000000110000011100000111000111110
(saturated image)	7	0000111000001110000011100000111000101110000011100001111000111110
(image with text)	7	0000111000001110000011100000111000101110000011100001111000111110
(original image)	7	0000111000001110000011100000111000101110000011100001111000111110
(a different cat image)	36	0000000001111100111111001111100011111000111110000111110000000000

Even with a distorted image, the engine picks up on similar images based on visual content.

It works well on slight image modifications like saturation, text content, etc. The distance is almost same for all the paper bag cat images, as the visual content in all those images are similar, and, ahash does a nice job by not drastically changing the hash content.

A different cat image, should have different visual content as compared to the input image, hence, a higher distance overall.

This implementation should be effective for small-scale usecases like blocking spam/restricted images based on a pre-built library of restricted/spam images, as it is resistant to smaller image modifications/augmentations like watermarks, text overlays, etc.

Optimizations ¶

This is a basic implementation using average hash as the core fingerprinting algorithm. While average hash (ahash) is a simple and efficient algorithm for generating perceptual hashes, there are some drawbacks:

ahash isn’t resistant to image rotation, scaling, and translation.
Since ahash is a simple algorithm, and is based on the average of the image pixels, it can be insensitive to local features, which impacts the accuracy of the search engine.

Other developments have been made in this space, and one such technique is difference hash (dHash), which computes the difference between adjacent pixels in an image. Other hashing algorithms include pHash, wavelet hash, etc. Using these techniques might improve the accuracy of the engine.

The current implementation works fine, scanning and comparing against each image is good enough for ~1k images. But it definitely won’t scale.

As the size of the database grows, the time it takes to compare the input image’s ahash to all the hashes in the database can become a bottleneck. BK trees could be a good fit in this space, as we’re dealing with distances.

A BK tree is a data structure that is used to store a set of items in a way that allows for efficient searching. For this project, a BK tree could be used to store the perceptual hashes of a set of images. The tree is constructed in such a way that images with similar perceptual hashes are stored close to each other in the tree.

To search for an image in the BK tree, we traverse the tree and compare the distance of the query image hash to the hashes in the nodes. This approach significantly reduces the number of hash comparisons that need to be performed, resulting in faster search times.

Conclusion ¶

This was a fun project overall; got to learn things like perceptual hashing and efficient searching techniques. The next step is to add multiple hashing algorithms like dHash, pHash and improve the search performance by using BK-trees.

You can check out the project at github.com/burntcarrot/hashsearch.

Go plugins

2022-07-07T00:00:00Z

Exploring Go's plugin package.

cobra is a library for creating powerful modern CLI applications, and is currently used by popular projects like Kubernetes, Hugo, etc.

This post focuses on adding plugins to a Cobra application, as an example for introducing Go’s plugin package.

Go and Plugins ¶

Plugins are useful for extending an application’s feature list. Go generally supports two kinds of plugins:

Compile-time plugins ¶

Compile-time plugins consist of code packages that get compiled into the application’s main binary. Compile-time plugins are easy to discover and register as they’re baked in the binary itself.

Compile-time plugins

Run-time plugins ¶

Run-time plugins hook up to an application at run-time.

A special build mode enables compiling packages into shared object (.so) libraries. The plugin package provides simple functions for loading shared libraries and getting symbols from them.

Run-time plugins

An example ¶

Let’s try using run-time plugins for our Cobra application. As a first step, we will need to build a CLI using Cobra:

Code can be found here.

// main.go

package main

import (
	"errors"
	"fmt"
	"log"
	"plugin"

	"github.com/spf13/cobra"
)

func main() {
	mainCmd := &cobra.Command{
		Use:   "main",
		Short: "Simple cobra app",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println("Main app called!")
		},
	}

	// Load the plugin.
	pluginCmd, err := LoadPlugin("./plugin/plugin.so", "PluginCmd")
	if err != nil {
		log.Fatalf("failed to load plugin: %v\n", err)
	}

	// Register the plugin at runtime.
	mainCmd.AddCommand(pluginCmd)

	if err := mainCmd.Execute(); err != nil {
		log.Fatalf("error: %v\n", err)
	}
}

The above code snippet does the following:

Create a main command
Load and register the plugin (explained later)
Execute the main command

Here is the code for loading a plugin:

// main.go

// LoadPlugin loads a cobra command from a shared object file.
func LoadPlugin(pluginPath, cmdName string) (*cobra.Command, error) {
	p, err := plugin.Open(pluginPath)
	if err != nil {
		return nil, err
	}

	c, err := p.Lookup("New" + cmdName)
	if err != nil {
		return nil, err
	}

	// Get access to the plugin function and get the command.
	pluginFunc, ok := c.(func() *cobra.Command)
    if !ok {
        return nil, errors.New("failed to perform a lookup.")
    }

	pluginCmd := pluginFunc()

	return pluginCmd, nil
}

In the above code snippet, we use Open() for loading the plugin object. We assume that the function exposing the command has this naming convention: New<CommandName>, for example, NewPluginCmd, etc.

Since there is no way for looking up all symbols (exported functions), having a naming convention helps, as it ensures that the lookups are easy to maintain.

Since we have the command name, we perform a lookup using our naming convention. Once we have the symbol, we assert the type of the symbol. In our case, we expect the signature of func() *cobra.Command. This function is used for retrieving the cobra command, which is eventually returned to the caller.

Let’s implement the plugin:

// plugin/plugin.go

package main

import (
	"fmt"

	"github.com/spf13/cobra"
)

func NewPluginCmd() *cobra.Command {
	pluginCmd := &cobra.Command{
		Use:   "plugin",
		Short: "Plugin command",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println("Plugin called!")
		},
	}

	return pluginCmd
}

We can now build the plugin using -buildmode=plugin. The plugin buildmode expects you to have at least one main package.

go build -buildmode=plugin .

Tada! We can see a plugin.so in our directory:

$ tree
.
├── go.mod
├── go.sum
├── main.go
└── plugin
    ├── plugin.go
    └── plugin.so

1 directory, 5 files

Here is a visualization on how the run-time plugin is registered into our Cobra app:

Registering a plugin

Let’s run the main app:

$ go run main.go --help
Simple cobra app

Usage:
  main [flags]
  main [command]

Available Commands:
  completion  Generate the autocompletion script for the specified shell
  help        Help about any command
  plugin      Plugin command

Flags:
  -h, --help   help for main

Use "main [command] --help" for more information about a command.

Nice! The plugin got registered to the main Cobra app at runtime!

Let’s run the plugin:

$ go run main.go plugin
Plugin called!

We now have a plugin loaded at run-time, and is accessible through the main CLI.

Challenges ¶

Run-time plugins are a nice addition; however, they have some challenges:

Plugins don’t work in Windows. (related: #19282.)
The main app and the plugin must be compiled with the same Go version and must have the same value for the GOPATH variable.
Packages imported by the main app and the plugin must have the same version.

Tips ¶

Here are some tips while building and using plugins:

Plugins should focus on providing one functionality only.
Developers who import plugins at run time must consider a plugin as a black box.
Always ensure that third-party plugins do not get access to critical components like databases, etc. Remember, the black box analogy? Third-party plugins aren’t trustworthy at all times.
Since plugins are independent components that are loaded at runtime, extensive documentation is required. Symbol names used for lookup should be documented properly.
Plugins shouldn’t be slow. Slower the plugin, it has high impact on the performance of the program, as they’re loaded at runtime.

Distributing plugins ¶

Ok, we have plugins. But how do we distribute them?

The easiest way to distribute plugins would be to use a plugin registry/index. The main app can request for a plugin from the registry and load the plugin at runtime.

Here is a neat visualization on how the distribution flow works:

Distributing plugins

The main app asks for plugin2 from the “Plugin Handler” component. The Plugin Handler component has two components:

Plugin Fetcher: This component fetches the plugin from the registry and stores it in the user’s filesystem.
Plugin Loader: Checks the central directory for downloaded plugins, returns the requested plugin.

The plugin fetcher downloads plugin2 from the registry, which is eventually returned by the plugin loader. The main app now loads plugin2 at runtime.

Conclusion ¶

Plugins allow an application to extend its functionality. Plugin support was requested by the Cobra community a couple of times (#691, #1026, #1361), but at the time of posting this blog, plugins are not natively supported by Cobra.

Plugins can be a great way to reduce your binary size. Your CLI can choose to have a subset of your available commands, and the commands with lesser priority can be set up at runtime. This however comes at a cost of maintenance, but allows builds to be faster due to less bloat in your CLI.

I would also recommend taking a look at RPC-based plugin systems, for example, hashicorp/go-plugin and natefinch/pie.