Linux and GLX example of fast glReadPixels with PBOs

standalone_readpixels.c

/*
 * Minimal Linux- and GLX-specific example of using a pool of Pixel Buffer
 * Objects to stream pixel data.
 *
 * This is a minimal test case for opening an OpenGL window with raw
 * X11/Xlib/GLX, drawing a simple test image, and fast PBO-based streaming of
 * image data to disk with glReadPixels().
 *
 * On an NVidia RTX2080 Super, on Ubuntu 20.04, saving to an SSD, this sustains
 * 150+fps (peaking at 500+, hitting 900+ if only writing a single file instead
 * of many) and is still fast even with just one PBO.  It filled up my entire
 * /tmp in about 5 seconds.
 *
 * Compile with:
 *
 * gcc -DGL_GLEXT_PROTOTYPES -DGLX_GLXEXT_PROTOTYPES -std=gnu99 -Wall -Wextra \
 *  -Werror -fPIC -rdynamic -pipe -O3 -march=native -ffast-math \
 *  standalone_readpixels.c -o standalone_readpixels \
 *  -ldl -lm -lpng -lGLU -lGL -lX11 -lXext
 *
 * Note that in proper usage you should use something like GLEW or libepoxy
 * instead of defining the *_PROTOTYPES macros.  Also you'd probably avoid
 * old-style OpenGL 1.x drawing.
 *
 * This file is released to the public domain, or under the CC0 license where
 * that is not possible, in July, 2020 by Mike Bourgeous, but attribution is
 * appreciated.
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <math.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <dlfcn.h>
#include <GL/gl.h>
#include <GL/glu.h>
#include <GL/glext.h>
#include <GL/glx.h>
#include <GL/glxext.h>

#define WIDTH 1024
#define HEIGHT 768
#define NUM_PBOS 4
#define SINGLE_FILE 1
#define VSYNC 0

void download_pbo(GLuint *pbos, int pbo_idx, size_t pix_bytes, int frameno)
{
	printf("Binding read pbo %d to save frame %d\n", pbo_idx, frameno);

	GLuint pbo_id = pbos[pbo_idx];
	glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo_id);

	void *buf = glMapBuffer(GL_PIXEL_PACK_BUFFER, GL_READ_ONLY);
	if(buf == NULL) {
		fprintf(stderr, "Error mapping pixel buffer for reading\n");
		abort();
	}

	char filename[1024];
	if(SINGLE_FILE) {
		snprintf(filename, 1024, "/tmp/img.bgra.data");
	} else {
		snprintf(filename, 1024, "/tmp/img%08d.bgra.data", frameno);
	}
	int fd = open(filename, O_WRONLY | O_CREAT | O_CLOEXEC, S_IRUSR | S_IWUSR);
	if(fd < 0) {
		perror("Error opening output file for writing");
		abort();
	}
	if(write(fd, buf, pix_bytes) != (ssize_t)pix_bytes) {
		perror("Error writing data to output file");
	}

	close(fd);

	glUnmapBuffer(GL_PIXEL_PACK_BUFFER);

	glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
}

int main(void)
{
	Display *display = XOpenDisplay(NULL);
	if(display == NULL) {
		fprintf(stderr, "Error opening X display\n");
		return -1;
	}

	int screen = DefaultScreen(display);

	int glx_attr_list[] = {
		GLX_RGBA,
		GLX_DOUBLEBUFFER,
		GLX_RED_SIZE, 8,
		GLX_GREEN_SIZE, 8,
		GLX_BLUE_SIZE, 8,
		GLX_DEPTH_SIZE, 16,
		None
	};

	XVisualInfo *visual = glXChooseVisual(display, screen, glx_attr_list);
	if(visual == NULL) {
		fprintf(stderr, "Error choosing GLX visual\n");
		return -1;
	}

	GLXContext context = glXCreateContext(display, visual, 0, True);
	if(context == NULL) {
		fprintf(stderr, "Error creating GLX context\n");
		return -1;
	}

	XSetWindowAttributes attr = {
		.colormap = XCreateColormap(display, RootWindow(display, screen), visual->visual, AllocNone),
		.event_mask = KeyPressMask,
	};

	int x = 0;
	int y = 0;
	unsigned int w = WIDTH;
	unsigned int h = HEIGHT;
	Window window = XCreateWindow(
			display,
			RootWindow(display, screen),
			x, y,
			w, h,
			0,
			visual->depth,
			InputOutput,
			visual->visual,
			CWColormap | CWEventMask,
			&attr
		     );

	XSizeHints hints = {
		.flags = PMinSize | PMaxSize | PBaseSize,
		.min_width = w,
		.min_height = h,
		.max_width = w,
		.max_height = h,
		.base_width = w,
		.base_height = h,
	};
	XSetStandardProperties(display, window, "glReadPixels Speed Test", "glReadPixels Speed Test", 0, NULL, 0, &hints);

	XMapRaised(display, window);

	Window root_ignored;
	unsigned int border, depth;
	XGetGeometry(display, window, &root_ignored, &x, &y, &w, &h, &border, &depth);
	printf("Window size is %u by %u, display depth is %u\n", w, h, depth);

	glXMakeCurrent(display, window, context);

	glXSwapIntervalEXT(display, window, !!VSYNC);

	glViewport(0, 0, w, h);
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (float)w / (float)h, 0.1, 10.0);
	glMatrixMode(GL_MODELVIEW);

	glClearColor(0.0, 0.0, 0.0, 1.0);
	glClearDepth(10.0);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	glXSwapBuffers(display, window);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	glXSwapBuffers(display, window);

	GLuint pbos[NUM_PBOS];
	int read_pbo = -NUM_PBOS + 1;
	int write_pbo = 0;
	int read_frames = 0;
	int saved_frames = 0;
	int pix_bytes = w * h * 4;

	glGenBuffers(NUM_PBOS, pbos);
	for(int i = 0; i < NUM_PBOS; i++) {
		glBindBuffer(GL_PIXEL_PACK_BUFFER, pbos[i]);
		glBufferData(GL_PIXEL_PACK_BUFFER, pix_bytes, NULL, GL_STATIC_READ);
	}
	glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);

	glDisable(GL_DEPTH_TEST);

	glLoadIdentity();
	glTranslatef(0.0, 0.0, -sqrtf(3.0) * 1.5);

	int running = 1;
	while(running) {
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		glRotatef(1, 0.0, 1.0, 0.0);

		glBegin(GL_LINE_LOOP);
		glColor4f(1.0, 0.0, 0.0, 1.0);
		glVertex3f(-1, -0.5, 0);
		glColor4f(0.0, 1.0, 0.0, 1.0);
		glVertex3f(1, -0.5, 0);

		glColor4f(0.0, 0.0, 1.0, 1.0);
		glVertex3f(0, -1, 0);
		glColor4f(0.0, 1.0, 1.0, 1.0);
		glVertex3f(0, 1, 0);

		glColor4f(1.0, 0.0, 1.0, 1.0);
		glVertex3f(0, -0.5, -1);
		glColor4f(1.0, 1.0, 0.0, 1.0);
		glVertex3f(0, -0.5, 1);
		glEnd();

		glXSwapBuffers(display, window);

		read_frames++;
		printf("Binding write pbo %d for frame %d\n", write_pbo, read_frames);
		glReadBuffer(GL_FRONT);
		glBindBuffer(GL_PIXEL_PACK_BUFFER, pbos[write_pbo]);
		printf("Calling read pixels\n");
		glReadPixels(0, 0, w, h, GL_BGRA, GL_UNSIGNED_BYTE, NULL);
		glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
		write_pbo = (write_pbo + 1) % NUM_PBOS;

		if(read_pbo >= 0) {
			saved_frames++;
			download_pbo(pbos, read_pbo, pix_bytes, saved_frames);
		}
		read_pbo = (read_pbo + 1) % NUM_PBOS;

		XEvent event;
		while(XPending(display) > 0) {
			XNextEvent(display, &event);
			// Exit when escape is pressed
			if(event.type == KeyPress && XLookupKeysym(&event.xkey, 0) == XK_Escape) {
				running = 0;
			}
		}
	}

	printf("Draining read pbos\n");
	while(read_pbo != write_pbo) {
		if(read_pbo >= 0) {
			saved_frames++;
			download_pbo(pbos, read_pbo, pix_bytes, saved_frames);
		}
		read_pbo = (read_pbo + 1) % NUM_PBOS;
	}

	XCloseDisplay(display);

	return 0;
}