2 * Copyright (c) 2013 Emil Mikulic <emikulic@gmail.com>
4 * Permission to use, copy, modify, and distribute this software for any
5 * purpose with or without fee is hereby granted, provided that the above
6 * copyright notice and this permission notice appear in all copies.
8 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
16 * [ http://www.openbsd.org/cgi-bin/cvsweb/src/share/misc/license.template ]
19 // buftee: like tee(1) but buffers in memory.
21 // Read from stdin into memory all the time, so that the writer doesn't
22 // block. Our memory usage is unbounded.
24 // Write to a number of file descriptors whenever we're able to.
25 // Because write() can take a long time even with O_NONBLOCK set on the
26 // fd, we have to do this from a separate kernel scheduling entity (we use
29 // We use a thread per output filedescriptor so that one slow output (e.g. disk)
30 // doesn't block another (e.g. stdout to a terminal) (or vice-versa if you use
31 // flow control on the terminal!)
33 #define _GNU_SOURCE // for clock_gettime()
35 #include <sys/queue.h>
36 #include <sys/select.h>
37 #include <sys/syscall.h>
52 #define READ_BUF_SIZE 4096
53 #define SLOW_NSEC 4000
55 // Asserted on receipt of SIGTERM, SIGINT.
56 static volatile int stopping
= 0; // Global.
58 static void sig_stopping(int _ignored_
__attribute__((__unused__
))) {
62 // Reference-counted buffer, contains data that was read in the main thread.
70 struct buf_queue_elem
{
72 STAILQ_ENTRY(buf_queue_elem
) entries
;
74 STAILQ_HEAD(buf_queue
, buf_queue_elem
); // struct buf_queue
76 // Context for a writer thread.
77 struct writer_thread
{
81 // Each writer has its own queue.
82 struct buf_queue queue
;
84 // Pointer to a shared mutex which protects all queues.
85 pthread_mutex_t
* mutex
;
87 // When a writer runs out of work to do, it sleeps on a shared condition.
90 STAILQ_ENTRY(writer_thread
) entries
;
93 // A list of writer threads.
94 STAILQ_HEAD(writer_thread_list
, writer_thread
); // struct writer_thread_list
96 static struct writer_thread
* logger
= NULL
; // Global.
97 static struct timespec logger_start
; // Global.
99 static struct buf
* alloc_buf(const char* const data
, const int len
) {
101 struct buf
* buf
= malloc(sizeof(*buf
));
102 buf
->data
= malloc((size_t)len
);
103 memcpy(buf
->data
, data
, (size_t)len
);
109 static void unref_buf(struct buf
* buf
) {
110 assert(buf
->refcount
> 0);
111 if ((--buf
->refcount
) == 0) {
117 static void enqueue(struct buf_queue
* restrict queue
,
118 struct buf
* restrict buf
) {
119 struct buf_queue_elem
* elem
= malloc(sizeof(*elem
));
122 STAILQ_INSERT_TAIL(queue
, elem
, entries
);
125 static struct buf
* dequeue(struct buf_queue
* const queue
) {
126 assert(!STAILQ_EMPTY(queue
));
127 struct buf_queue_elem
* head
;
128 head
= STAILQ_FIRST(queue
);
129 STAILQ_REMOVE_HEAD(queue
, entries
);
130 struct buf
* buf
= head
->buf
;
135 static void xpthread_create(pthread_t
* thread
,
136 void* (*start_routine
)(void*),
138 int ret
= pthread_create(thread
, NULL
, start_routine
, arg
);
139 if (ret
== 0) return;
141 err(1, "pthread_create(%p) failed", thread
);
144 static void set_nonblocking(const int fd
) {
146 if ((flags
= fcntl(fd
, F_GETFL
)) == -1)
147 err(1, "fcntl(fd = %d, F_GETFL) failed", fd
);
148 if (fcntl(fd
, F_SETFL
, flags
| O_NONBLOCK
) == -1)
149 err(1, "fcntl(fd = %d, F_SETFL, O_NONBLOCK) failed", fd
);
152 static int make_file(const char* filename
) {
153 int fd
= open(filename
, O_CREAT
| O_EXCL
| O_NONBLOCK
| O_WRONLY
, 0666);
155 err(1, "failed to open(\"%s\")", filename
);
159 static void get_mono_time(struct timespec
* t
) {
160 if (clock_gettime(CLOCK_MONOTONIC
, t
) == -1)
161 err(1, "clock_gettime(CLOCK_MONOTONIC) failed");
164 static void time_diff(const struct timespec
* restrict start
,
165 const struct timespec
* restrict end
,
166 struct timespec
* restrict out
) {
167 out
->tv_sec
= end
->tv_sec
- start
->tv_sec
;
168 out
->tv_nsec
= end
->tv_nsec
- start
->tv_nsec
;
169 if (out
->tv_nsec
< 0) {
171 out
->tv_nsec
+= 1000000000;
173 assert(out
->tv_sec
>= 0);
174 assert(out
->tv_nsec
>= 0);
175 assert(out
->tv_nsec
< 1000000000);
178 static int gettid(void) {
179 return (int)syscall(SYS_gettid
);
182 static void nblogx(const char* format
, ...); // Forward.
184 static void warn_time(const char* desc
,
185 const struct timespec
* restrict start
,
186 const struct timespec
* restrict end
) {
187 if (pthread_self() == logger
->thread
) {
190 struct timespec diff
;
191 time_diff(start
, end
, &diff
);
192 if (diff
.tv_sec
> 0 || diff
.tv_nsec
> SLOW_NSEC
) {
193 nblogx("%s took %d.%09d secs", desc
, (int)diff
.tv_sec
, (int)diff
.tv_nsec
);
197 static int xread(const int fd
, char* const restrict buf
, const int count
) {
204 read_ret
= read(fd
, buf
, (size_t)count
);
207 warn_time("read()", &t0
, &t1
);
211 err(1, "read(fd = %d, count = %d) failed", fd
, count
);
214 assert(read_ret
>= 0);
215 return (int)read_ret
;
218 static void wait_until_writable(const int fd
) {
221 FD_SET(fd
, &write_fds
);
222 int select_ret
= select(fd
+ 1, NULL
, &write_fds
, NULL
, NULL
);
223 if (select_ret
== -1) {
224 if (errno
== EINTR
) {
225 assert(stopping
); // that should have been SIGTERM
228 err(1, "select(write fd = %d) failed", fd
);
230 if (!FD_ISSET(fd
, &write_fds
))
231 errx(1, "select() did not return writable fd = %d", fd
);
234 static int xwrite(const int fd
, struct buf
* const buf
) {
242 write_ret
= write(fd
, buf
->data
, (size_t)buf
->len
);
245 warn_time("write()", &t0
, &t1
);
248 if (write_ret
== -1) {
249 if (errno
== EAGAIN
) {
250 nblogx("write(fd = %d) got EAGAIN, sleeping and retrying", fd
);
251 wait_until_writable(fd
);
254 err(1, "write(fd = %d, count = %d) failed", fd
, buf
->len
);
258 assert(write_ret
>= 0);
259 if (write_ret
< buf
->len
)
260 err(1, "write(fd = %d, count = %d) stopped short (returned %d)",
261 fd
, buf
->len
, (int)write_ret
);
262 // FIXME: handle this
263 assert(write_ret
== buf
->len
);
264 return (int)write_ret
;
268 static void wait_until_readable(const int fd
) {
271 FD_SET(fd
, &read_fds
);
272 int select_ret
= select(fd
+ 1, &read_fds
, NULL
, NULL
, NULL
);
273 if (select_ret
== -1) {
274 if (errno
== EINTR
) {
275 assert(stopping
); // that should have been SIGTERM
278 err(1, "select(read fd = %d) failed", fd
);
280 if (!FD_ISSET(fd
, &read_fds
))
281 errx(1, "select() did not return readable fd = %d", fd
);
284 static void lock(pthread_mutex_t
* mutex
) {
285 int ret
= pthread_mutex_lock(mutex
);
286 if (ret
== 0) return;
288 err(1, "pthread_mutex_lock(%p) failed", mutex
);
291 static void unlock(pthread_mutex_t
* mutex
) {
292 int ret
= pthread_mutex_unlock(mutex
);
293 if (ret
== 0) return;
295 err(1, "pthread_mutex_unlock(%p) failed", mutex
);
298 static void* writer_routine(void *arg
) {
299 struct writer_thread
* my
= arg
;
300 struct buf
* buf
= NULL
;
303 while (!stopping
&& STAILQ_EMPTY(&my
->queue
)) {
305 pthread_cond_wait(my
->cond
, my
->mutex
);
307 if (!STAILQ_EMPTY(&my
->queue
))
308 buf
= dequeue(&my
->queue
);
315 int write_ret
= xwrite(my
->fd
, buf
);
316 if (write_ret
== 0) {
317 errx(1, "fd %d hit EOF", my
->fd
);
319 assert(write_ret
== buf
->len
);
321 // Unreference buffer, freeing it if we have to.
325 nblogx("thread exiting cleanly");
329 static void add_writer_thread(struct writer_thread_list
* list
,
331 pthread_mutex_t
* shared_queue_mutex
,
332 pthread_cond_t
* shared_wakeup_cond
) {
333 struct writer_thread
* writer
= malloc(sizeof(*writer
));
336 writer
->mutex
= shared_queue_mutex
;
337 writer
->cond
= shared_wakeup_cond
;
338 STAILQ_INIT(&writer
->queue
);
339 STAILQ_INSERT_TAIL(list
, writer
, entries
);
340 xpthread_create(&writer
->thread
, writer_routine
, writer
);
343 static void xpthread_cond_broadcast(pthread_cond_t
* cond
) {
344 int ret
= pthread_cond_broadcast(cond
);
345 if (ret
== 0) return;
347 err(1, "pthread_cond_broadcast(%p) failed", cond
);
350 static void init_logger_thread(void) {
351 assert(logger
== NULL
);
352 get_mono_time(&logger_start
);
353 logger
= malloc(sizeof(*logger
));
354 logger
->fd
= STDERR_FILENO
;
355 logger
->mutex
= malloc(sizeof(*logger
->mutex
));
356 pthread_mutex_init(logger
->mutex
, NULL
);
357 logger
->cond
= malloc(sizeof(*logger
->cond
));
358 pthread_cond_init(logger
->cond
, NULL
);
359 STAILQ_INIT(&logger
->queue
);
360 xpthread_create(&logger
->thread
, writer_routine
, logger
);
363 static void _nblog_helper(int want_errno
,
369 struct timespec diff
;
371 time_diff(&logger_start
, &now
, &diff
);
376 extern char *__progname
; // This is where glibc stashes argv[0].
377 len
= snprintf(buf
, sizeof(buf
), "%s: tid %d at %d.%09d: ",
378 __progname
, gettid(), (int)diff
.tv_sec
, (int)diff
.tv_nsec
);
381 len
+= vsnprintf(buf
+ len
, sizeof(buf
) - len
, format
, va
);
384 len
+= snprintf(buf
+ len
, sizeof(buf
) - len
, ": %s (errno = %d)",
385 strerror(saved_errno
), saved_errno
);
388 len
+= snprintf(buf
+ len
, sizeof(buf
) - len
, "\n");
389 struct buf
* b
= alloc_buf(buf
, len
);
393 enqueue(&logger
->queue
, b
);
394 xpthread_cond_broadcast(logger
->cond
);
395 unlock(logger
->mutex
);
398 // nblog() is like warn() but non-blocking.
399 static void nblog(const char* format
, ...) {
400 int saved_errno
= errno
;
402 va_start(va
, format
);
403 _nblog_helper(1, saved_errno
, format
, va
);
407 static void nblogx(const char* format
, ...) {
409 va_start(va
, format
);
410 _nblog_helper(0, 0, format
, va
);
414 static void xpthread_join(pthread_t thread
) {
415 int ret
= pthread_join(thread
, NULL
);
416 if (ret
== 0) return;
418 err(1, "pthread_join(%lu) failed", thread
);
421 static void sig_continue(int _ignored_
__attribute__((__unused__
))) {
422 set_nonblocking(STDERR_FILENO
);
425 int main(int argc
, char **argv
) {
426 init_logger_thread();
429 pthread_mutex_t shared_queue_mutex
= PTHREAD_MUTEX_INITIALIZER
;
430 pthread_cond_t shared_wakeup_cond
= PTHREAD_COND_INITIALIZER
;
432 struct writer_thread_list writers
;
433 STAILQ_INIT(&writers
);
435 if (signal(SIGINT
, sig_stopping
) == SIG_ERR
) err(1, "signal() failed");
436 if (signal(SIGTERM
, sig_stopping
) == SIG_ERR
) err(1, "signal() failed");
437 if (signal(SIGCONT
, sig_continue
) == SIG_ERR
) err(1, "signal() failed");
438 //if (signal(SIGPIPE, SIG_IGN) == SIG_ERR) err(1, "signal() failed");
441 // On Linux, making STDOUT non-blocking has the side-effect of
442 // also making STDIN nonblocking.
443 add_writer_thread(&writers
,
446 &shared_wakeup_cond
);
448 // Process cmdline args.
449 for (int i
= 1; i
< argc
; i
++) {
450 add_writer_thread(&writers
,
453 &shared_wakeup_cond
);
458 wait_until_readable(STDIN_FILENO
);
460 nblogx("stopping after select()");
465 char data
[READ_BUF_SIZE
];
466 int read_ret
= xread(STDIN_FILENO
, data
, sizeof(data
));
468 nblogx("stdin hit EOF");
471 struct buf
* buf
= alloc_buf(data
, read_ret
);
474 lock(&shared_queue_mutex
);
475 struct writer_thread
* writer
;
476 STAILQ_FOREACH(writer
, &writers
, entries
) {
477 enqueue(&(writer
->queue
), buf
);
479 xpthread_cond_broadcast(&shared_wakeup_cond
);
480 unlock(&shared_queue_mutex
);
483 // Wake and join threads.
484 lock(&shared_queue_mutex
);
486 xpthread_cond_broadcast(&shared_wakeup_cond
);
487 unlock(&shared_queue_mutex
);
489 struct writer_thread
* writer
;
490 STAILQ_FOREACH(writer
, &writers
, entries
) {
491 xpthread_join(writer
->thread
);
496 while (!STAILQ_EMPTY(&writers
)) {
497 struct writer_thread
* writer
= STAILQ_FIRST(&writers
);
498 STAILQ_REMOVE_HEAD(&writers
, entries
);
499 // FIXME: free its queue?
503 // Clean up logger thread.
504 nblogx("writer threads stopped");
506 xpthread_cond_broadcast(logger
->cond
);
507 unlock(logger
->mutex
);
508 xpthread_join(logger
->thread
);
511 warnx("exiting cleanly");
514 // vim:set ts=2 sw=2 tw=80 et: