n-way-output-port.rkt

#lang racket/base

(require (only-in racket/port open-output-nowhere)
         (only-in racket/string string-join))

(struct pending-write (bstr out start end) #:authentic #:mutable)

(define-syntax-rule (define/pending-write (bstr out start end) pending-write)
  (define-values (bstr out start end)
    (let ([x pending-write])
      (values (pending-write-bstr x)
              (pending-write-out x)
              (pending-write-start x)
              (pending-write-end x)))))

(define (n-way-output-port . output-ports)
  (define-values (out0 out-rest)
    (if (null? output-ports)
        (values (open-output-nowhere 'nowhere #f) null)
        (values (car output-ports) (cdr output-ports))))

  (define pending null)
  (define lock (make-semaphore 1))

  (define ready-evt
    (for/fold ([evt out0]) ([out (in-list out-rest)])
      (replace-evt out (λ (_) evt))))
  
  (define retry-evt
    (wrap-evt ready-evt (λ (_) #f)))

  (define port-name
    (string->symbol
     (format "n-way-output (~a)"
             (string-join (map (λ (port) (format "~s" (object-name port)))
                               (cons out0 out-rest))
                          " "))))

  (define (write-pending!)
    (when (pair? pending)
      (let loop ([xs pending])
        (cond
          [(null? xs)
           (set! pending null)]
          [else
           (define/pending-write (bstr out start end) (car xs))
           (define n (or (write-bytes-avail* bstr out start end) 0))
               
           (cond
             [(= n (- end start))
              (loop (cdr xs))]
             [else
              (when (> n 0)
                (set-pending-write-start! (car xs) (+ start n)))
              (unless (eq? pending xs)
                (set! pending xs))])]))))

  (define (flush-all)
    (flush-output out0)
    
    (for ([out (in-list out-rest)])
      (flush-output out)))
  
  (define (write-out bstr start end non-blocking? enable-break?)
    (define result
      (call-with-semaphore
       lock
       (λ ()
         (write-pending!)

         (cond [(pair? pending) retry-evt]
               [(= start end)   0]
               [enable-break?   (write-out* write-bytes-avail/enable-break bstr start end)]
               [else            (write-out* write-bytes-avail* bstr start end)]))
       (λ () retry-evt)))

    (when (eqv? result 0)
      (flush-all))

    result)

  (define (write-out* write-initial bstr start end)
    (define m (write-initial bstr out0 start end))

    (cond
      [(or (not m) (= m 0))
       retry-evt]
      [else
       (let loop ([xs out-rest])
         (when (pair? xs)
           (define n (or (write-bytes-avail* bstr (car xs) start (+ start m)) 0))

           (cond
             [(= n m)
              (loop (cdr xs))]
             [else
              (set! pending
                    (cons (pending-write bstr (car xs) (+ start n) (+ start m))
                          (map (λ (out-i)
                                 (pending-write bstr out-i 0 (+ start m)))
                               (cdr xs))))])))
       m]))

  (define (close)
    (call-with-semaphore
     lock
     (λ ()
       (for ([x (in-list pending)])
         (define/pending-write (bstr out start end) x)
         (write-bytes bstr out start end))
       (set! pending null)))

    (flush-all))

  (make-output-port
   port-name
   ready-evt
   write-out
   close))