;;; Nar Herder
;;;
;;; Copyright © 2021 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 storage)
  #:use-module (srfi srfi-1)
  #:use-module (ice-9 ftw)
  #:use-module (ice-9 match)
  #:use-module (ice-9 threads)
  #:use-module (rnrs bytevectors)
  #:use-module (web uri)
  #:use-module (web client)
  #:use-module (web response)
  #:use-module (logging logger)
  #:use-module (logging port-log)
  #:use-module (prometheus)
  #:use-module (json)
  #:use-module ((guix build utils) #:select (dump-port mkdir-p))
  #:use-module ((guix store) #:select (store-path-hash-part))
  #:use-module (nar-herder utils)
  #:use-module (nar-herder database)
  #:export (store-item-in-local-storage?
            remove-nar-files-by-hash

            get-nar-files

            start-nar-removal-thread
            start-mirroring-thread))

(define (store-item-in-local-storage? database storage-root hash)
  (let ((narinfo-files (database-select-narinfo-files database hash)))
    (when (null? narinfo-files)
      (error "no narinfo files"))
    (every (lambda (file)
             (file-exists?
              (string-append storage-root
                             (uri-decode
                              (assq-ref file 'url)))))
           narinfo-files)))

(define (remove-nar-files-by-hash database storage-root metrics-registry
                                  hash)
  (define nar-files-metric
    (or (metrics-registry-fetch-metric metrics-registry
                                       "nar_files_total")
        (make-gauge-metric metrics-registry
                           "nar_files_total"
                           #:labels '(stored))))

  (let ((narinfo-files (database-select-narinfo-files database hash)))
    (when (null? narinfo-files)
      (error "no narinfo files"))
    (for-each
     (lambda (file)
       (let* ((filename
               (string-append storage-root
                              (uri-decode
                               (assq-ref file 'url))))
              (exists?
               (file-exists? filename)))
         (when exists?
           (remove-nar-from-storage storage-root
                                    (assq-ref file 'url)))

         (metric-decrement nar-files-metric
                           #:label-values
                           `((stored . ,(if exists? "true" "false"))))))
     narinfo-files)))

(define (get-storage-size storage-root)
  (define enter? (const #t))
  (define (leaf name stat result)
    (+ result
       (or (and=> (stat:blocks stat)
                  (lambda (blocks)
                    (* blocks 512)))
           (stat:size stat))))

  (define (down name stat result) result)
  (define (up name stat result) result)
  (define (skip name stat result) result)

  (define (error name stat errno result)
    (format (current-error-port) "warning: ~a: ~a~%"
            name (strerror errno))
    result)

  (file-system-fold enter? leaf down up skip error
                           0 ; Start counting at 0
                           storage-root))

(define (remove-nar-from-storage storage-root nar-file)
  (let* ((filename
          (string-append storage-root "/" nar-file)))
    (log-msg 'INFO "removing nar " nar-file)
    (delete-file filename))
  #t)

(define (index-storage database storage-root)
  (define (get-files-hash)
    (define storage-root-length
      (string-length storage-root))

    (define enter? (const #t))
    (define (leaf name stat result)
      (hash-set! result
                 (string-drop name storage-root-length)
                 #t)
      result)

    (define (down name stat result) result)
    (define (up name stat result) result)
    (define (skip name stat result) result)

    (define (error name stat errno result)
      (format (current-error-port) "warning: ~a: ~a~%"
              name (strerror errno))
      result)

    (file-system-fold enter? leaf down up skip error
                      (make-hash-table (expt 2 19))
                      storage-root))

  (let* ((files-hash
          (if storage-root
              (get-files-hash)
              (make-hash-table)))
         (narinfo-files
          (database-map-all-narinfo-files
           database
           (lambda (file)
             (let* ((url (uri-decode
                          (assq-ref file 'url)))
                    (stored? (hash-ref files-hash url)))
               (when stored?
                 ;; Delete the hash entry, so
                 ;; that the hash at the end will
                 ;; just contain the unrecognised
                 ;; files
                 (hash-remove! files-hash url))

               `(,@file
                 (stored? . ,stored?)))))))
    `((narinfo-files      . ,narinfo-files)
      (unrecognised-files . ,(hash-map->list (lambda (key _) key)
                                             files-hash)))))

(define* (get-nar-files database storage-root metrics-registry
                        #:key stored?)
  (define nar-files-metric
    (or (metrics-registry-fetch-metric metrics-registry
                                       "nar_files_total")
        (make-gauge-metric metrics-registry
                           "nar_files_total"
                           #:labels '(stored))))

  (let* ((index (index-storage database storage-root))
         (selected-files
          (filter
           (lambda (file)
             (eq? (assq-ref file 'stored?) stored?))
           (assq-ref index 'narinfo-files))))

    (let ((selected-files-count
           (length selected-files))
          (all-files-count
           (length (assq-ref index 'narinfo-files))))

      (metric-set nar-files-metric
                  selected-files-count
                  #:label-values `((stored . ,(if stored? "true" "false"))))
      (metric-set nar-files-metric
                  (- all-files-count selected-files-count)
                  #:label-values `((stored . ,(if stored? "false" "true")))))

    selected-files))

(define (start-nar-removal-thread database
                                  storage-root storage-limit
                                  metrics-registry
                                  nar-removal-criteria)
  (define storage-size-metric
    (make-gauge-metric metrics-registry
                       "storage_size_bytes"))

  (define nar-files-metric
    (or (metrics-registry-fetch-metric metrics-registry
                                       "nar_files_total")
        (make-gauge-metric metrics-registry
                           "nar_files_total"
                           #:labels '(stored))))

  (define (check-removal-criteria nar criteria)
    (define narinfo
      (database-select-narinfo-for-file database (assq-ref nar 'url)))

    (match criteria
      (('and and-criteria)
       (every (lambda (c)
                (check-removal-criteria nar c))
              and-criteria))
      (('stored-on url)
       (let ((uri (string->uri
                   (string-append (if (symbol? url)
                                      (symbol->string url)
                                      url)
                                  "/"
                                  (store-path-hash-part
                                   (assq-ref narinfo 'store-path))
                                  ".narinfo/info"))))
         (call-with-values
             (lambda ()
               (retry-on-error
                (lambda ()
                  (http-get uri #:decode-body? #f))
                #:times 3
                #:delay 5))
           (lambda (response body)
             (and (= (response-code response)
                     200)

                  (let ((json-body (json-string->scm
                                    (utf8->string body))))
                    (eq? (assoc-ref json-body "stored")
                         #t)))))))))

  (define (nar-can-be-removed? nar)
    (any (lambda (criteria)
           (check-removal-criteria nar criteria))
         nar-removal-criteria))

  (define (run-removal-pass)
    (log-msg 'INFO "looking for nars to remove")
    (let ((initial-storage-size
           (with-time-logging "getting storage size"
             (get-storage-size storage-root))))
      (log-msg 'DEBUG "initial storage size " initial-storage-size)
      (metric-set storage-size-metric
                  initial-storage-size)

      (let loop ((storage-size
                  initial-storage-size)
                 (stored-nar-files
                  (with-time-logging "getting stored nar files"
                    (get-nar-files database storage-root metrics-registry
                                   #:stored? #t))))
        ;; Look through items in local storage, check if the removal
        ;; criteria have been met, and if so, delete it

        (unless (null? stored-nar-files)
          (let ((nar-to-consider (car stored-nar-files)))
            (if (nar-can-be-removed? nar-to-consider)
                (begin
                  (remove-nar-from-storage
                   storage-root
                   (uri-decode
                    (assq-ref nar-to-consider 'url)))

                  (metric-decrement nar-files-metric
                                    #:label-values '((stored . "true")))
                  (metric-increment nar-files-metric
                                    #:label-values '((stored . "false")))

                  (let ((storage-size-estimate
                         (- storage-size
                            (assq-ref nar-to-consider 'size))))
                    (when (> storage-size-estimate storage-limit)
                      (loop storage-size-estimate
                            (cdr stored-nar-files)))))
                (loop storage-size
                      (cdr stored-nar-files)))))))
    (log-msg 'INFO "finished looking for nars to remove"))

  (when (null? nar-removal-criteria)
    (error "must be some removal criteria"))

  (call-with-new-thread
   (lambda ()
     (while #t
       (with-exception-handler
           (lambda (exn)
             (log-msg 'ERROR "nar removal pass failed " exn))
         run-removal-pass
         #:unwind? #t)

       (sleep 300)))))

(define (start-mirroring-thread database mirror storage-limit storage-root
                                metrics-registry)

  (define no-storage-limit?
    (not (integer? storage-limit)))

  (define storage-size-metric
    (make-gauge-metric metrics-registry
                       "storage_size_bytes"))

  (define nar-files-metric
    (or (metrics-registry-fetch-metric metrics-registry
                                       "nar_files_total")
        (make-gauge-metric metrics-registry
                           "nar_files_total"
                           #:labels '(stored))))

  (define (fetch-file file)
    (let* ((string-url
            (string-append mirror file))
           (uri
            (string->uri (string-append mirror file)))
           (destination-file-name
            (string-append storage-root
                           (uri-decode file)))
           (tmp-file-name
            (string-append destination-file-name "-tmp")))
      (log-msg 'INFO "fetching " string-url)

      (mkdir-p (dirname destination-file-name))

      (when (file-exists? tmp-file-name)
        (delete-file tmp-file-name))

      (call-with-values
          (lambda ()
            (http-get uri
                      #:decode-body? #f
                      #:streaming? #t))
        (lambda (response body)
          (unless (= (response-code response)
                     200)
            (error "unknown response code"))

          (call-with-output-file tmp-file-name
            (lambda (output-port)
              (dump-port body output-port)))
          (rename-file tmp-file-name
                       destination-file-name)

          (metric-increment nar-files-metric
                            #:label-values '((stored . "true")))
          (metric-decrement nar-files-metric
                            #:label-values '((stored . "false")))))))

  (define (download-nars initial-storage-size)
    ;; If there's free space, then consider downloading missing nars
    (when (< initial-storage-size storage-limit)
      (let loop ((storage-size initial-storage-size)
                 (missing-nar-files (get-nar-files
                                     database storage-root metrics-registry
                                     #:stored? #f)))
        (unless (null? missing-nar-files)
          (let ((file (car missing-nar-files)))
            (log-msg 'DEBUG "considering "
                     (assq-ref file 'url))
            (let ((file-bytes (assq-ref file 'size)))
              (if (or no-storage-limit?
                      (< (+ storage-size file-bytes)
                         storage-limit))
                  (let ((success?
                         (with-exception-handler
                             (lambda (exn)
                               (log-msg 'ERROR "failed to fetch "
                                        (assq-ref file 'url)
                                        ": " exn)
                               #f)
                           (lambda ()
                             (fetch-file (assq-ref file 'url))
                             #t)
                           #:unwind? #t)))
                    (loop (if success?
                              (+ storage-size file-bytes)
                              storage-size)
                          (cdr missing-nar-files)))
                  ;; This file won't fit, so try the next one
                  (loop storage-size
                        (cdr missing-nar-files)))))))))

  (define (fast-download-nars)
    (define parallelism 3)

    (let ((missing-nar-files (get-nar-files
                              database storage-root metrics-registry
                              #:stored? #f)))
      (n-par-for-each
       parallelism
       (lambda (file)
         (log-msg 'DEBUG "considering "
                  (assq-ref file 'url))
         (with-exception-handler
             (lambda (exn)
               (log-msg 'ERROR "failed to fetch "
                        (assq-ref file 'url)
                        ": " exn)
               #f)
           (lambda ()
             (fetch-file (assq-ref file 'url))
             #t)
           #:unwind? #t))
       missing-nar-files)))

  (define (run-mirror-pass)
    (log-msg 'DEBUG "running mirror pass")
    (let ((initial-storage-size (with-time-logging "getting storage size"
                                  (get-storage-size storage-root))))
      (metric-set storage-size-metric
                  initial-storage-size)
      (if no-storage-limit?
          (fast-download-nars)
          (download-nars initial-storage-size)))
    (log-msg 'DEBUG "finished mirror pass"))

  (call-with-new-thread
   (lambda ()
     (while #t
       (with-exception-handler
           (lambda (exn)
             (log-msg 'ERROR "mirror pass failed " exn))
         run-mirror-pass
         #:unwind? #t)

       (sleep 300)))))