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 err := mergeMmdb(); err != nil { log.Err(err).Msg("Failed to merge mmdb") rollBackMmdb(backup, backupErr) 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, backupErr) 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. func rollBackMmdb(backup string, backupErr error) { if backupErr != nil { log.Err(backupErr).Msg("No mmdb backup available; cannot restore") } else 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") } }