path: root/nar-herder/cached-compression.scm

;;; Nar Herder
;;;
;;; Copyright © 2022, 2023 Christopher Baines <mail@cbaines.net>
;;;
;;; This program is free software: you can redistribute it and/or
;;; modify it under the terms of the GNU Affero General Public License
;;; as published by the Free Software Foundation, either version 3 of
;;; the License, or (at your option) any later version.
;;;
;;; This program 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
;;; Affero General Public License for more details.
;;;
;;; You should have received a copy of the GNU Affero General Public
;;; License along with this program.  If not, see
;;; <http://www.gnu.org/licenses/>.

(define-module (nar-herder cached-compression)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (srfi srfi-71)
  #:use-module (ice-9 ftw)
  #:use-module (ice-9 match)
  #:use-module (ice-9 atomic)
  #:use-module (ice-9 threads)
  #:use-module (logging logger)
  #:use-module (prometheus)
  #:use-module (web uri)
  #:use-module (web client)
  #:use-module (web response)
  #:use-module (guix store)
  #:use-module ((guix utils) #:select (compressed-file?
                                       call-with-decompressed-port))
  #:use-module ((guix build utils)
                #:select (dump-port mkdir-p))
  #:use-module (nar-herder utils)
  #:use-module (nar-herder database)
  #:export (make-maybe-trigger-creation-of-compressed-nars))

;; Nar caching overview
;;
;; On start
;;  - Compute the size of each cached compression directory
;;  - Remove database entries if they're missing from the directory
;;  - Remove files if they're missing from the database
;;  - Remove (random TODO) files if the size exceeds the max
;;
;; On nar usage
;;  - Bump count
;;  - If count is sufficient
;;    - Trigger generation of the cached nar
;;      - Skip if the work queue already includes this job
;;      - At the start of the job, check for a database entry and exit
;;        early if one exists
;;    - Update the size of the cached compression directory
;;      - If the size exceeds the max, remove (random TODO) files

(define (perform-cached-compression-startup database
                                            enabled-cached-compressions)
  (let* ((database-entries-missing-files
          ;; List tracking entries in the database for cached files,
          ;; where the file is missing from the disk.
          ;;
          ;; These database entries will be removed at the end of the
          ;; startup procecss.
          '())

         ;; alist of compression to hash table of files
         ;;
         ;; Entries from the hash tables will be removed when the
         ;; database entry is processed below, so these hash tables
         ;; will be left reflecting files in the directories, but with
         ;; no entry in the database.
         ;;
         ;; These files will be deleted at the end of the startup
         ;; process.
         (files-by-compression
          (map
           (match-lambda
             ((compression . details)
              (let ((result (make-hash-table)))
                (for-each
                 (lambda (file)
                   (hash-set! result file #t))
                 (scandir (assq-ref details 'directory)
                          (negate (cut member <> '("." "..")))))
                (cons compression
                      result))))
           enabled-cached-compressions))

         (cached-bytes-by-compression
          (database-fold-cached-narinfo-files
           database
           (lambda (details result)
             (let ((compression
                    (assq-ref details 'compression))
                   (filename
                    (store-path-base
                     (assq-ref details 'store-path))))

               (let ((files-hash
                      (assq-ref files-by-compression compression)))
                 (if (hash-ref files-hash filename)
                     (begin
                       (hash-remove! files-hash filename)

                       `((,compression . ,(+ (assq-ref details 'size)
                                             (or (assq-ref result compression)
                                                 0)))
                         ,@(alist-delete compression result)))

                     ;; Database entry, but file missing
                     (begin
                       (set! database-entries-missing-files
                             (cons details
                                   database-entries-missing-files))
                       result)))))
           (map
            (lambda (compression)
              (cons compression 0))
            (map car enabled-cached-compressions)))))

    ;; Delete cached files with no database entries
    (for-each
     (match-lambda
       ((compression . hash-table)
        (let ((count (hash-count (const #t)
                                 hash-table)))
          (unless (= 0 count)
            (let ((directory
                   (assq-ref
                    (assq-ref enabled-cached-compressions compression)
                    'directory)))
              (simple-format #t "deleting ~A cached files from ~A\n"
                             count
                             directory)
              (hash-for-each
               (lambda (filename _)
                 (delete-file (string-append directory "/" filename)))
               hash-table))))))
     files-by-compression)

    ;; Delete database entries where the file is missing
    (let ((count (length database-entries-missing-files)))
      (unless (= 0 count)
        (simple-format
         #t
         "deleting ~A cached_narinfo_files entries due to missing files\n"
         count)

        (for-each
         (lambda (details)
           (database-remove-cached-narinfo-file database
                                                (assq-ref details 'narinfo-id)
                                                (symbol->string
                                                 (assq-ref details 'compression))))
         database-entries-missing-files)))

    cached-bytes-by-compression))

(define (maybe-remove-cached-files-for-compression database
                                                   enabled-cached-compressions
                                                   compression
                                                   cached-bytes-by-compression-box
                                                   nar-cache-bytes-metric)
  (let* ((compression-details
          (assq-ref enabled-cached-compressions
                    compression))
         (max-size
          (assq-ref compression-details
                    'directory-max-size)))
    (when max-size
      (let ((current-size
             (assq-ref
              (atomic-box-ref cached-bytes-by-compression-box)
              compression)))
        (when (> current-size max-size)
          (let ((bytes-to-remove
                 (- current-size max-size)))
            (log-msg 'DEBUG
                     "looking to remove " bytes-to-remove " bytes of "
                     compression " compressed nars")

            (match
                (database-fold-cached-narinfo-files
                 database
                 (lambda (details result)
                   (match result
                     (#(finished? bytes files-to-remove)
                      (if finished?
                          result
                          (if (eq? (assq-ref details 'compression)
                                   compression)
                              (let ((new-bytes
                                     (+ bytes
                                        (assq-ref details 'size))))
                                ;; finished if enough bytes are going
                                ;; to be removed
                                (vector (>= new-bytes bytes-to-remove)
                                        new-bytes
                                        (cons details files-to-remove)))
                              result)))))
                 #(#f 0 ()))
              (#(#t bytes-for-removal files-to-remove-details)
               (log-msg 'DEBUG "removing " (length files-to-remove-details)
                        " " compression " compressed nars from the cache")

               ;; Use an explicit transaction as it handles the
               ;; database being busy,
               (database-call-with-transaction
                database
                (lambda _
                  (for-each
                   (lambda (details)
                     ;; Remove all the database entries first, as
                     ;; that'll stop these files appearing in narinfos
                     (database-remove-cached-narinfo-file
                      database
                      (assq-ref details 'narinfo-id)
                      (symbol->string compression)))
                   files-to-remove-details)))

               (let ((directory
                      (assq-ref compression-details 'directory)))
                 (for-each
                  (lambda (details)
                    (let ((filename
                           (string-append
                            directory "/"
                            (basename (assq-ref details 'store-path)))))
                      (log-msg 'DEBUG "deleting " filename)
                      (delete-file filename)))
                  files-to-remove-details)

                 (let* ((cached-bytes-by-compression
                         (atomic-box-ref cached-bytes-by-compression-box))
                        (new-cached-bytes
                         (- (assq-ref cached-bytes-by-compression
                                      compression)
                            bytes-for-removal)))
                   (atomic-box-set!
                    cached-bytes-by-compression-box
                    (alist-cons
                     compression
                     new-cached-bytes
                     (alist-delete
                      compression
                      cached-bytes-by-compression)))

                   (metric-set
                    nar-cache-bytes-metric
                    new-cached-bytes
                    #:label-values `((compression . ,compression)))

                   (log-msg 'DEBUG "finished removing " bytes-for-removal
                            " bytes of " compression
                            " cached nars from " directory
                            " (new size " new-cached-bytes ")")))))))))))

(define* (make-maybe-trigger-creation-of-compressed-nars
          database
          metrics-registry
          nar-source
          enabled-cached-compressions
          cached-compression-min-uses
          #:key (cached-compression-workers 2) scheduler)

  (define nar-cache-bytes-metric
    (make-gauge-metric metrics-registry
                       "nar_cache_size_bytes"
                       #:labels '(compression)))

  (let ((consider-nar-request-channel
         (make-worker-thread-set
          (lambda ()
            (let ((cached-bytes-by-compression-box
                   (make-atomic-box #f)))
              (atomic-box-set!
               cached-bytes-by-compression-box
               (perform-cached-compression-startup database
                                                   enabled-cached-compressions))

              (for-each
               (match-lambda
                 ((compression . bytes)
                  (metric-set
                   nar-cache-bytes-metric
                   bytes
                   #:label-values `((compression . ,compression)))))
               (atomic-box-ref cached-bytes-by-compression-box))

              ;; Remove cached files if the max size is exceeded
              (for-each
               (match-lambda
                 ((compression . _)
                  (maybe-remove-cached-files-for-compression
                   database
                   enabled-cached-compressions
                   compression
                   cached-bytes-by-compression-box
                   nar-cache-bytes-metric)))
               (atomic-box-ref cached-bytes-by-compression-box))

              (list cached-bytes-by-compression-box)))
          #:name "comp nar req"
          ;; Just make one thread, as this thread won't do much work
          ;; and relies on a hash table that shouldn't be accessed by
          ;; multiple threads
          #:parallelism 1)))

    (define with-usage-hash-table
      (let ((nar-cached-compression-usage-hash-table
             (make-hash-table 65536)))
        (lambda (proc)
          (monitor
           (proc nar-cached-compression-usage-hash-table)))))

    (define (proc narinfo-id cached-bytes-by-compression-box)
      (let* ((existing-cached-files
              ;; This is important both to avoid trying to create
              ;; files twice, but also in the case where additional
              ;; compressions are enabled and more files need to be
              ;; generated
              (database-select-cached-narinfo-files-by-narinfo-id
               database
               narinfo-id))
             (existing-compressions
              (map (lambda (details)
                     (assq-ref details 'compression))
                   existing-cached-files))
             (missing-compressions
              (lset-difference eq?
                               (map car enabled-cached-compressions)
                               existing-compressions))
             (compress-file?
              (let ((url
                     (assq-ref
                      (first
                       (database-select-narinfo-files-by-narinfo-id
                        database
                        narinfo-id))
                      'url)))

                ;; TODO: Maybe this should be configurable?
                (not (compressed-file? url)))))

        (when (and compress-file?
                   (not (null? missing-compressions)))
          (let ((new-count
                 (with-usage-hash-table
                  (lambda (usage)
                    (let ((val (+ 1
                                  (or (hash-ref usage narinfo-id)
                                      0))))
                      (hash-set! usage
                                 narinfo-id
                                 val)
                      val)))))

            (when (> new-count
                     cached-compression-min-uses)
              (for-each
               (lambda (missing-compression)
                 (let ((new-bytes
                        (make-compressed-nar
                         database
                         nar-source
                         enabled-cached-compressions
                         narinfo-id
                         missing-compression
                         #:level (assq-ref
                                  (assq-ref enabled-cached-compressions
                                            missing-compression)
                                  'level))))

                   ;; Do this here, after creating the new cached nar,
                   ;; just in case there's a lack of space
                   (monitor
                    (let* ((cached-bytes-by-compression
                            (atomic-box-ref
                             cached-bytes-by-compression-box))
                           (updated-bytes-for-compression
                            (+ (assq-ref cached-bytes-by-compression
                                         missing-compression)
                               new-bytes)))

                      (metric-set
                       nar-cache-bytes-metric
                       updated-bytes-for-compression
                       #:label-values `((compression . ,missing-compression)))

                      (atomic-box-set!
                       cached-bytes-by-compression-box
                       (alist-cons
                        missing-compression
                        updated-bytes-for-compression
                        (alist-delete
                         missing-compression
                         cached-bytes-by-compression)))

                      (maybe-remove-cached-files-for-compression
                       database
                       enabled-cached-compressions
                       missing-compression
                       cached-bytes-by-compression-box
                       nar-cache-bytes-metric)))))
               missing-compressions)

              (with-usage-hash-table
               (lambda (usage)
                 (hash-remove! usage narinfo-id))))))))

    (let ((process-job
           count-jobs
           count-threads
           list-jobs
           (create-work-queue cached-compression-workers
                              proc
                              #:name "cached compression")))

      (lambda (narinfo-id)
        (spawn-fiber
         (lambda ()
           (call-with-worker-thread
            consider-nar-request-channel
            (lambda (cached-bytes-by-compression-box)
              (let ((in-progress-narinfo-ids
                     (map car (list-jobs))))

                (unless (member narinfo-id in-progress-narinfo-ids)
                  (process-job narinfo-id cached-bytes-by-compression-box)))
              #t)))
         scheduler
         #:parallel? #t)))))

(define* (make-compressed-nar database
                              nar-source
                              enabled-cached-compressions
                              narinfo-id
                              target-compression
                              #:key level)
  (define cached-compression-details
    (assq-ref enabled-cached-compressions target-compression))

  (define narinfo-files
    (database-select-narinfo-files-by-narinfo-id database
                                                 narinfo-id))

  (log-msg 'DEBUG "making " target-compression " for "
           (uri-decode
            (basename
             (assq-ref (first narinfo-files) 'url))))

  (let* ((source-narinfo-file
          ;; There's no specific logic to this, it should be possible
          ;; to use any file
          (first narinfo-files))
         (source-filename
          (cond
           ((string-prefix? "http" nar-source)
            (let* ((output-port (mkstemp "/tmp/nar-herder-source-nar-XXXXXX"))
                   (filename
                    (port-filename output-port))
                   (uri
                    (string->uri
                     (string-append nar-source
                                    (assq-ref source-narinfo-file 'url)))))

              (log-msg 'DEBUG "downloading " (uri->string uri))
              (with-exception-handler
                  (lambda (exn)
                    (close-port output-port)
                    (delete-file filename)
                    (raise-exception exn))
                (lambda ()
                  (with-fibers-port-timeouts
                   (lambda ()
                     (call-with-values
                         (lambda ()
                           (let ((port
                                  socket
                                  (open-socket-for-uri* uri)))
                             (http-get uri
                                       #:port port
                                       #:decode-body? #f
                                       #:streaming? #t)))
                       (lambda (response body)
                         (unless (= (response-code response)
                                    200)
                           (error "unknown response code"
                                  (response-code response)))

                         (dump-port body output-port)))
                     (close-port output-port))
                   #:timeout 30))
                #:unwind? #t)

              filename))
           ((string-prefix? "/" nar-source)
            (string-append
             ;; If it's a filename, then it's the canonical path to
             ;; the storage directory
             nar-source
             (uri-decode (assq-ref source-narinfo-file 'url))))
           (else
            (error "unknown nar source")))))

    (let* ((dest-directory
            (assq-ref cached-compression-details
                      'directory))
           (dest-filename
            (string-append
             dest-directory
             "/"
             (uri-decode
              (basename
               (assq-ref source-narinfo-file 'url)))))
           (tmp-dest-filename
            (string-append dest-filename ".tmp")))

      (when (file-exists? tmp-dest-filename)
        (delete-file tmp-dest-filename))
      (when (file-exists? dest-filename)
        (delete-file dest-filename))

      (mkdir-p dest-directory)

      (call-with-input-file
          source-filename
        (lambda (source-port)
          (call-with-decompressed-port
              (string->symbol
               (assq-ref source-narinfo-file
                         'compression))
              source-port
            (lambda (decompressed-source-port)
              (let ((call-with-compressed-output-port*
                     (match target-compression
                       ('gzip
                        (@ (zlib) call-with-gzip-output-port))
                       ('lzip
                        (@ (lzlib) call-with-lzip-output-port))
                       ('zstd
                        (@ (zstd) call-with-zstd-output-port))
                       ('none
                        (lambda (port proc)
                          (proc port))))))
                (apply
                 call-with-compressed-output-port*
                 (open-output-file tmp-dest-filename)
                 (lambda (compressed-port)
                   (dump-port decompressed-source-port
                              compressed-port))
                 (if level
                     `(#:level ,level)
                     '())))))))
      (rename-file
       tmp-dest-filename
       dest-filename)

      (when (string-prefix? "http" nar-source)
        (log-msg 'DEBUG "deleting temporary file " source-filename)
        (delete-file source-filename))

      (let ((bytes
             (stat:size (stat dest-filename))))

        (database-call-with-transaction
         database
         (lambda _
           (database-insert-cached-narinfo-file
            database
            narinfo-id
            bytes
            (symbol->string target-compression))))

        (log-msg 'DEBUG "created " dest-filename)

        bytes))))