ip-info/internal/services/db_updater/main.go

package db_updater

import (
	"github.com/robfig/cron/v3"
	"github.com/rs/zerolog/log"
	"os"
	"path"
	"sync/atomic"
	"thuanle.me/ip-info/configs"
	"thuanle.me/ip-info/internal/data"
	"thuanle.me/ip-info/pkg/osx"
)

// DbUpdatedAt stores the DB last-updated time as unix milliseconds.
// Accessed concurrently by the cron goroutine (writer) and the /metrics
// handler (reader), so it must be synchronized — hence atomic.Int64.
var DbUpdatedAt atomic.Int64

// DbUpdatedAtMillis returns the last DB update time in unix milliseconds,
// or 0 if the DB has never been updated.
func DbUpdatedAtMillis() int64 {
	return DbUpdatedAt.Load()
}

func StartUpdateDbService() {
	c := cron.New()
	_, _ = c.AddFunc("@daily", fetchDbs)
	c.Start()
	fetchDbs()
}

func fetchDbs() {
	newFlag := false
	for _, url := range configs.GeoDbSourcePaths {
		newFlag = download(url) || newFlag
	}

	if newFlag {
		log.Info().Msg("New DB downloaded. Recreating mmdb")
		// mmdbmeld.WriteMMDB truncates the canonical mmdb before writing. If
		// the merge fails the file would be left empty/corrupt — and because
		// download() already persisted the etag, the next daily run would get
		// 304 Not Modified and never retry, leaving the service broken until a
		// restart (which then fails to open the empty file). Back the file up
		// first so we can restore it on failure.
		backup, backupErr := backupMmdb()
		if backupErr != nil {
			// We could not back up the current canonical mmdb. Do NOT proceed:
			// mergeMmdb() truncates the canonical file before writing, and with
			// no backup a later merge/reload failure would leave it corrupt
			// (unrecoverable on restart). Keep serving the current DB and let
			// the next run retry.
			log.Err(backupErr).Msg("Failed to back up mmdb, skipping rebuild")
			return
		}
		if err := mergeMmdb(); err != nil {
			log.Err(err).Msg("Failed to merge mmdb")
			rollBackMmdb(backup)
			return
		}

		err := data.Ins().Reload()
		if err != nil {
			log.Err(err).Msg("Failed to reload mmdb")
			// Reload failed on the newly merged file: restore the previous
			// canonical mmdb so we keep serving known-good data.
			rollBackMmdb(backup)
			return
		}

		// Both merge and reload succeeded: the new canonical file is live, the
		// backup is no longer needed.
		cleanupMmdbBackup(backup)

		// Stamp the update time from the freshly written canonical file.
		stampDbUpdatedAt()
		return
	}

	// No new download: on startup this is the path that loads an existing DB.
	// Ensure the timestamp reflects the file we are actually serving.
	if DbUpdatedAt.Load() == 0 {
		stampDbUpdatedAt()
	}
}

// stampDbUpdatedAt records the mtime of the canonical mmdb (unix ms) as the
// last DB update time. This covers both initial load (pre-existing files) and
// post-download reloads, so /metrics never reports 0 while serving real data.
func stampDbUpdatedAt() {
	info, err := os.Stat(configs.MmdbDbFile)
	if err != nil {
		log.Err(err).Str("file", configs.MmdbDbFile).Msg("Failed to stat mmdb")
		return
	}
	DbUpdatedAt.Store(info.ModTime().UnixMilli())
}

// backupMmdb copies the canonical mmdb to a sibling backup file. It returns
// ("", nil) when there is no canonical file yet (first-ever build), in which
// case there is nothing to restore on failure.
func backupMmdb() (string, error) {
	if _, err := os.Stat(configs.MmdbDbFile); os.IsNotExist(err) {
		return "", nil
	}
	backup := configs.MmdbDbFile + ".bak"
	if err := osx.Copy(configs.MmdbDbFile, backup); err != nil {
		return "", err
	}
	return backup, nil
}

// rollBackMmdb restores the canonical mmdb from backup (if any) and deletes
// the etag files so the next cron run re-downloads and retries the merge.
// backup is "" on a first-ever build (no canonical file existed to back up),
// in which case there is nothing to restore.
func rollBackMmdb(backup string) {
	if backup != "" {
		if err := osx.Copy(backup, configs.MmdbDbFile); err != nil {
			log.Err(err).Msg("Failed to restore mmdb backup")
		} else {
			log.Info().Str("file", configs.MmdbDbFile).Msg("Restored mmdb from backup")
		}
		_ = os.Remove(backup)
	}
	// Drop etags so a failed merge is retried on the next run instead of being
	// masked by a 304 Not Modified.
	for _, url := range configs.GeoDbSourcePaths {
		etag := configs.GeoDbFolder + path.Base(url) + ".etag"
		if err := os.Remove(etag); err != nil && !os.IsNotExist(err) {
			log.Err(err).Str("file", etag).Msg("Failed to remove etag")
		}
	}
}

// cleanupMmdbBackup removes a successful merge's backup file.
func cleanupMmdbBackup(backup string) {
	if backup == "" {
		return
	}
	if err := os.Remove(backup); err != nil && !os.IsNotExist(err) {
		log.Err(err).Str("file", backup).Msg("Failed to remove mmdb backup")
	}
}