;;; remote-worker.scm -- Remote build worker.
;;; Copyright © 2020 Mathieu Othacehe <othacehe@gnu.org>
;;;
;;; This file is part of Cuirass.
;;;
;;; GNU Guix is free software; you can redistribute it and/or modify it
;;; under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 3 of the License, or (at
;;; your option) any later version.
;;;
;;; GNU Guix is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with GNU Guix.  If not, see <http://www.gnu.org/licenses/>.

(define-module (cuirass remote-worker)
  #:use-module (cuirass remote)
  #:use-module (gcrypt pk-crypto)
  #:use-module (guix)
  #:use-module (guix avahi)
  #:use-module (guix config)
  #:use-module (guix diagnostics)
  #:use-module (guix pki)
  #:use-module (guix records)
  #:use-module (guix scripts)
  #:use-module (guix ui)
  #:use-module (guix build syscalls)
  #:use-module (guix scripts publish)
  #:use-module (simple-zmq)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (srfi srfi-34)
  #:use-module (srfi srfi-37)
  #:use-module (ice-9 atomic)
  #:use-module (ice-9 match)
  #:use-module (ice-9 threads)

  #:export (remote-worker))

;; Indicate if the process has to be stopped.
(define %stop-process?
  (make-atomic-box #f))

(define (show-help)
  (format #t (G_ "Usage: remote-worker [OPTION]...
Start a remote build worker.\n"))
  (display (G_ "
  -w, --workers=COUNT       start COUNT parallel workers"))
  (display (G_ "
  -p, --publish-port=PORT   publish substitutes on PORT"))
  (display (G_ "
      --public-key=FILE     use FILE as the public key for signatures"))
  (display (G_ "
      --private-key=FILE    use FILE as the private key for signatures"))
  (newline)
  (display (G_ "
  -h, --help                display this help and exit"))
  (display (G_ "
  -V, --version             display version information and exit"))
  (newline)
  (show-bug-report-information))

(define %options
  (list (option '(#\h "help") #f #f
                (lambda _
                  (show-help)
                  (exit 0)))
        (option '(#\V "version") #f #f
                (lambda _
                  (show-version-and-exit "guix publish")))
        (option '(#\w "workers") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'workers (string->number* arg) result)))
        (option '(#\p "publish-port") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'publish-port (string->number* arg) result)))
        (option '("public-key") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'public-key-file arg result)))
        (option '("private-key") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'private-key-file arg result)))))

(define %default-options
  `((workers . 1)
    (publish-port . 5558)
    (public-key-file . ,%public-key-file)
    (private-key-file . ,%private-key-file)))

\f
;;;
;;; ZMQ connection.
;;;

(define %zmq-context
  (zmq-create-context))

(define (zmq-backend-endpoint address port)
  "Return a ZMQ endpoint identifying the build server available by TCP at
ADDRESS and PORT."
  (string-append "tcp://" address ":" (number->string port)))

(define (zmq-dealer-socket)
  "The ZMQ socket to communicate with the worker threads."
  (zmq-create-socket %zmq-context ZMQ_DEALER))

\f
;;;
;;; Worker.
;;;

;; The port of the local publish server.
(define %local-publish-port
  (make-atomic-box #f))

(define (server-publish-url address port)
  "Return the server publish url at ADDRESS and PORT."
  (string-append "http://" address ":" (number->string port)))

(define (service-txt->publish-port txt)
  "Parse the service TXT record and return the server publish port."
  (define (parse-txt)
    (fold (lambda (param params)
            (match (string-split param #\=)
              ((key value)
               (cons (cons (string->symbol key) value)
                     params))))
          '()
          txt))

  (let ((params (parse-txt)))
    (string->number (assq-ref params 'publish))))

(define (service->publish-url service)
  "Return the URL of the publish server corresponding to the service with the
given NAME."
  (let* ((address (avahi-service-address service))
         (txt (avahi-service-txt service))
         (publish-port
          (service-txt->publish-port txt)))
    (server-publish-url address publish-port)))

(define (service->local-publish-url service)
  "Return the URL of the local publish server."
  (let* ((local-address (avahi-service-local-address service))
         (port (atomic-box-ref %local-publish-port)))
    (server-publish-url local-address port)))

(define* (run-build drv service #:key reply)
  "Build DRV and send messages upon build start, failure or completion to the
build server identified by SERVICE-NAME using the REPLY procedure.

The publish server of the build server is added to the list of the store
substitutes-urls.  This way derivations that are not present on the worker can
still be substituted."
  (with-store store
    (let ((publish-url (service->publish-url service))
          (local-publish-url (service->local-publish-url service)))
      (add-substitute-url store publish-url)
      (reply (zmq-build-started-message drv))
      (guard (c ((store-protocol-error? c)
                 (info (G_ "Derivation `~a' build failed: ~a~%")
                       drv (store-protocol-error-message c))
                 (reply (zmq-build-failed-message drv))))
        (if (build-derivations store (list drv))
            (reply (zmq-build-succeeded-message drv local-publish-url))
            (reply (zmq-build-failed-message drv)))))))

(define* (run-command command service #:key reply)
  "Run COMMAND.  SERVICE-NAME is the name of the build server that sent the
command.  REPLY is a procedure that can be used to reply to this server."
  (match (zmq-read-message command)
    (('build ('drv drv) ('system system))
     (info (G_ "Building `~a' derivation.~%") drv)
     (run-build drv service #:reply reply))
    (('no-build)
     #t)))

(define (start-worker worker service)
  "Start a worker thread named NAME, reading commands from the DEALER socket
and executing them.  The worker can reply on the same socket."
  (define (reply socket client)
    (lambda (message)
      (zmq-send-msg-parts-bytevector
       socket
       (list (zmq-empty-delimiter) client
             (zmq-empty-delimiter) (string->bv message)))))

  (define (ready socket)
    (zmq-send-msg-parts-bytevector
     socket
     (list (make-bytevector 0)
           (string->bv
            (zmq-worker-ready-message (worker->sexp worker))))))

  (define (request-work socket)
    (let ((name (worker-name worker)))
      (zmq-send-msg-parts-bytevector
       socket
       (list (make-bytevector 0)
             (string->bv (zmq-worker-request-work-message name))))))

  (call-with-new-thread
   (lambda ()
     (set-thread-name (worker-name worker))
     (let* ((socket (zmq-dealer-socket))
            (address (avahi-service-address service))
            (port (avahi-service-port service))
            (endpoint (zmq-backend-endpoint address port)))
       (zmq-connect socket endpoint)
       (ready socket)
       (let loop ()
         (request-work socket)
         (match (zmq-get-msg-parts-bytevector socket '())
           ((empty client empty command)
            (run-command (bv->string command) service
                         #:reply (reply socket client))))
         (sleep 1)
         (loop))))))

\f
;;;
;;; Entry point.
;;;

;; The PID of the publish process.
(define %publish-pid
  (make-atomic-box #f))

(define (signal-handler)
  "Catch SIGINT to stop the Avahi event loop and the publish process before
exiting."
  (sigaction SIGINT
    (lambda (signum)
      (let ((publish-pid (atomic-box-ref %publish-pid)))
        (atomic-box-set! %stop-process? #t)

        (and publish-pid
             (begin
               (kill publish-pid SIGHUP)
               (waitpid publish-pid)))

        (exit 1)))))

(define (remote-worker args)
  (with-error-handling
    (let* ((opts (args-fold* args %options
                             (lambda (opt name arg result)
                               (leave (G_ "~A: unrecognized option~%") name))
                             (lambda (arg result)
                               (leave (G_ "~A: extraneous argument~%") arg))
                             %default-options))
           (workers (assoc-ref opts 'workers))
           (publish-port (assoc-ref opts 'publish-port))
           (public-key
            (read-file-sexp
             (assoc-ref opts 'public-key-file)))
           (private-key
            (read-file-sexp
             (assoc-ref opts 'private-key-file))))

      (atomic-box-set! %local-publish-port publish-port)

      (atomic-box-set!
       %publish-pid
       (publish-server publish-port
                       #:public-key public-key
                       #:private-key private-key))

      (avahi-browse-service-thread
       (lambda (action service)
         (case action
           ((new-service)
            (for-each (lambda (n)
                        (start-worker (worker
                                       (name (generate-worker-name))
                                       (systems '("x86_64-linux")))
                                      service))
                      (iota workers)))))
       #:types (list remote-server-service-type)
       #:stop-loop? (lambda ()
                      (atomic-box-ref %stop-process?))))))