cuticle/hi/x264enc.c

#include"node.h"

#include<stdio.h>
#include<assert.h>
#include<stdlib.h>
#include<string.h>

#include<sys/random.h>
#include<unistd.h>
#include<fcntl.h>
#include<sys/types.h>
#include<sys/stat.h>
#include<errno.h>
#include<pthread.h>

#include"mode.h"
#include"img.h"

typedef struct {
	CHiPubNode pub;
	
	char fifoIn[256];
	
	int fdOut;
	int fdIn;

	pthread_t iothread;
	pthread_mutex_t iomutex;
	
	size_t outQueueLen;
	uint8_t *outQueue;
	
	size_t inQueueLen;
	size_t inQueueCap;
	uint8_t *inQueue;
	
	size_t waitingForNum;
} Internal;

void *iothread_func(void *ud) {
	Internal *n = ud;
	
	while(1) {
		pthread_mutex_lock(&n->iomutex);
		
		while(1) {
			uint8_t toread[4096];
			ssize_t readcount = read(n->fdIn, toread, sizeof(toread));
			
			if(readcount == -1) {
				if(errno != EAGAIN && errno != EWOULDBLOCK) {
					goto error;
				} else {
					break;
				}
			}
			
			if(n->inQueueLen + readcount > n->inQueueCap) {
				n->inQueue = realloc(n->inQueue, n->inQueueCap += 8192);
			}
			
			memcpy(n->inQueue + n->inQueueLen, toread, readcount);
			n->inQueueLen += readcount;
			
			if(readcount < sizeof(toread)) {
				break;
			}
		}
		
		//printf("inqueue %lu outqueue %lu\n\n", n->inQueueLen, n->outQueueLen);
		
		while(n->outQueueLen > 0) {
			ssize_t wrotecount = write(n->fdOut, n->outQueue, n->outQueueLen);
			
			if(wrotecount == -1) {
				if(errno != EAGAIN && errno != EWOULDBLOCK) {
					goto error;
				} else {
					break;
				}
			}
			
			memmove(n->outQueue, n->outQueue + wrotecount, n->outQueueLen - wrotecount);
			n->outQueueLen -= wrotecount;
			
			if(wrotecount < n->outQueueLen) {
				break;
			}
		}
		
		pthread_mutex_unlock(&n->iomutex);
	}
	
	return NULL;
error:
	pthread_mutex_unlock(&n->iomutex);
	return NULL;
}

int encodeh264_start(CHiPubNode *pub) {
	Internal *n = (Internal*) pub;
	
	CHiImage *firstFrame = CHi_Crawl(&pub->sinks[0])->data.sample;
	
#ifndef _WIN32
	uint8_t randoms[4];
	getrandom(randoms, sizeof(randoms), 0);
	
	snprintf(n->fifoIn, sizeof(n->fifoIn), "/tmp/cuticlex264_%02X%02X%02X%02X.264", randoms[0], randoms[1], randoms[2], randoms[3]);
	
	char res[256];
	snprintf(res, sizeof(res), "%ix%i", firstFrame->width, firstFrame->height);
	
	assert(mkfifo(n->fifoIn, 0666) == 0);
	
	int fd[2];
	pipe(fd);
	
	if(fork() == 0) {
		close(fd[1]);
		dup2(fd[0], STDIN_FILENO);
		close(fd[0]);
		
		if(execlp("x264", "x264", "--profile", "main", "--tune", "zerolatency", "--preset", "ultrafast", "--input-csp", "bgra", "--input-depth", "16", "--input-res", res, "--aud", "--demuxer", "raw", "--fps", "30", "--qp", "50", "-o", n->fifoIn, "-", NULL) == -1) {
			abort();
		}
	} else {
		close(fd[0]);
		n->fdOut = fd[1];
	}
	
	n->fdIn = open(n->fifoIn, O_RDONLY | O_NONBLOCK, 0);
	
	fcntl(n->fdOut, F_SETFL, fcntl(n->fdOut, F_GETFL) | O_NONBLOCK);
#endif
	
	n->outQueueLen = 0;
	n->outQueue = NULL;
	
	n->inQueueLen = 0;
	n->inQueueCap = 8192;
	n->inQueue = malloc(sizeof(*n->inQueue) * n->inQueueCap);
	
	n->waitingForNum = 0;
	
	assert(pthread_create(&n->iothread, NULL, iothread_func, n) == 0);
	
	return 1;
}

int encodeh264_stop(CHiPubNode *pub) {
	Internal *n = (Internal*) pub;
	
	close(n->fdOut);
	close(n->fdIn);
	
	unlink(n->fifoIn);
	
	return 1;
}

static uint8_t *find_aud(uint8_t *bs, size_t len, int *sz) {
	if(len <= 3) {
		return NULL;
	}
	
	for(size_t i = 3; i < len; i++) {
		if(bs[i] == 9 && bs[i - 1] == 1 && bs[i - 2] == 0 && bs[i - 3] == 0) {
			if(i >= 4 && bs[i - 4] == 0) {
				*sz = 4;
				return &bs[i - 4];
			} else {
				*sz = 3;
				return &bs[i - 3];
			}
		}
	}
	
	return NULL;
}

int encodeh264_perform(CHiPubNode *pub) {
	Internal *n = (Internal*) pub;
	
	pthread_mutex_lock(&n->iomutex);
	{
		//if(CHi_GetMode() == CUTIHI_MODE_OFFLINE || n->waitingForNum == 0) {
			CHiImage *frame = CHi_Crawl(&pub->sinks[0])->data.sample;
			
			size_t frameSize = frame->width * 8 * frame->height;
			
			n->outQueue = realloc(n->outQueue, n->outQueueLen + frameSize);
			CHi_Restride(frame->data8, n->outQueue + n->outQueueLen, frame->stride, frame->width * 8, frame->height);
			n->outQueueLen += frameSize;
			
		//	n->waitingForNum++;
		//}
	}
	
	CHiBSFrames *frames = calloc(1, sizeof(*frames));
	
	while(1) {
		int aud0sz;
		uint8_t *aud0 = find_aud(n->inQueue, n->inQueueLen, &aud0sz);
		
		if(aud0 == NULL) {
			break;
		}
		
		// First AUD must always be at the start because it's the delimiter
		assert(aud0 == n->inQueue);
		
		int aud1sz;
		uint8_t *aud1 = find_aud(n->inQueue + aud0sz, n->inQueueLen - aud0sz, &aud1sz);
		
		// Second AUD must exist, otherwise we don't actually know if the packet's over
		if(aud1 == NULL) {
			break;
		}
		
		size_t framesz = aud1 - aud0;
		
		static int nextts = 0;
		nextts++;
		
		frames = realloc(frames, sizeof(*frames) + sizeof(CHiBSFrame) * (frames->count + 1));
		//frames->data[frames->count].timestamp = CHi_Time_Get(pub->ng) * 1000;
		frames->data[frames->count].timestamp = nextts * 33;
		frames->data[frames->count].sz = framesz - 6;
		memcpy(frames->data[frames->count].ptr = malloc(framesz - 6), aud0 + 6, framesz - 6);
		frames->data[frames->count].flags = 0;
		frames->count++;
		
		memmove(n->inQueue, n->inQueue + framesz, n->inQueueLen - framesz);
		n->inQueueLen -= framesz;
		
		//if(n->waitingForNum > 0) {
		//	n->waitingForNum--;
		//}
	}
	
	pthread_mutex_unlock(&n->iomutex);
	
	pub->sources[0].type = CUTIHI_VAL_H264BS;
	pub->sources[0].data.bitstream = frames;
	
	return 1;
}

CUTIVIS CHiPubNode *CHi_EncodeH264() {
	Internal *ret = calloc(1, sizeof(*ret));
	ret->pub.type = CUTIHI_T('CEnc', 'H264');
	ret->pub.Start = encodeh264_start;
	ret->pub.Perform = encodeh264_perform;
	ret->pub.Stop = encodeh264_stop;
	ret->pub.sinks = calloc(sizeof(*ret->pub.sinks), ret->pub.sinkCount = 1);
	ret->pub.sources = calloc(sizeof(*ret->pub.sources), ret->pub.sourceCount = 1);
	return &ret->pub;
}