3 files changed, 349 insertions, 0 deletions
diff --git a/db/cron.go b/db/cron.go
new file mode 100644
index 0000000..a7282a4
--- /dev/null
+++ b/db/cron.go
@@ -0,0 +1,139 @@
+package db
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/benbusby/farside/services"
+	"github.com/robfig/cron/v3"
+)
+
+const defaultPrimary = "https://farside.link/state"
+const defaultCFPrimary = "https://cf.farside.link/state"
+
+var LastUpdate time.Time
+
+func InitCronTasks() {
+	log.Println("Initializing cron tasks...")
+
+	cronDisabled := os.Getenv("FARSIDE_CRON")
+	if len(cronDisabled) == 0 || cronDisabled == "1" {
+		c := cron.New()
+		c.AddFunc("@every 10m", queryServiceInstances)
+		c.AddFunc("@daily", updateServiceList)
+		c.Start()
+	}
+
+	queryServiceInstances()
+}
+
+func updateServiceList() {
+	fileName := services.GetServicesFileName()
+	_, _ = services.FetchServicesFile(fileName)
+	services.InitializeServices()
+}
+
+func queryServiceInstances() {
+	log.Println("Starting instance queries...")
+
+	isPrimary := os.Getenv("FARSIDE_PRIMARY")
+	if len(isPrimary) == 0 || isPrimary != "1" {
+		remoteServices, err := fetchInstancesFromPrimary()
+		if err != nil {
+			log.Println("Unable to fetch instances from primary", err)
+		}
+
+		for _, service := range remoteServices {
+			SetInstances(service.Type, service.Instances)
+		}
+
+		return
+	}
+
+	for _, service := range services.ServiceList {
+		fmt.Printf("===== %s =====\n", service.Type)
+		var instances []string
+		for _, instance := range service.Instances {
+			testURL := strings.ReplaceAll(
+				service.TestURL,
+				"<%=query%>",
+				"current+weather")
+			available := queryServiceInstance(
+				instance,
+				testURL,
+			)
+
+			if available {
+				instances = append(instances, instance)
+			}
+		}
+
+		SetInstances(service.Type, instances)
+	}
+
+	LastUpdate = time.Now().UTC()
+}
+
+func fetchInstancesFromPrimary() ([]services.Service, error) {
+	primaryURL := defaultPrimary
+	useCF := os.Getenv("FARSIDE_CF_ENABLED")
+	if len(useCF) > 0 && useCF == "1" {
+		primaryURL = defaultCFPrimary
+	}
+
+	resp, err := http.Get(primaryURL)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	bodyBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	var serviceList []services.Service
+	err = json.Unmarshal(bodyBytes, &serviceList)
+	return serviceList, err
+}
+
+func queryServiceInstance(instance, testURL string) bool {
+	testMode := os.Getenv("FARSIDE_TEST")
+	if len(testMode) > 0 && testMode == "1" {
+		return true
+	}
+
+	ua := "Mozilla/5.0 (compatible; Farside/1.0.0; +https://farside.link)"
+	url := instance + testURL
+
+	req, err := http.NewRequest(http.MethodGet, url, nil)
+	if err != nil {
+		fmt.Println("    [ERRO] Failed to create new http request!", err)
+		return false
+	}
+
+	req.Header.Set("User-Agent", ua)
+	client := &http.Client{
+		Timeout: 10 * time.Second,
+	}
+	resp, err := client.Do(req)
+
+	if err != nil {
+		fmt.Println("    [ERRO] Error fetching instance:", err)
+		return false
+	} else if resp.StatusCode != http.StatusOK {
+		fmt.Printf("    [WARN] Received non-200 status for %s\n", url)
+		return false
+	} else {
+		fmt.Printf("    [INFO] Received 200 status for %s\n", url)
+	}
+
+	return true
+}
diff --git a/db/db.go b/db/db.go
new file mode 100644
index 0000000..efdf145
--- /dev/null
+++ b/db/db.go
@@ -0,0 +1,150 @@
+package db
+
+import (
+	"encoding/json"
+	"errors"
+	"log"
+	"math/rand"
+	"os"
+	"slices"
+	"time"
+
+	"github.com/benbusby/farside/services"
+	"github.com/dgraph-io/badger/v4"
+)
+
+var (
+	badgerDB     *badger.DB
+	selectionMap map[string]string
+
+	cachedServiceList []services.Service
+	cacheUpdated      time.Time
+)
+
+func InitializeDB() error {
+	var err error
+
+	dbDir := os.Getenv("FARSIDE_DB_DIR")
+	if len(dbDir) == 0 {
+		dbDir = "./badger-db"
+	}
+
+	badgerDB, err = badger.Open(badger.DefaultOptions(dbDir))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func SetInstances(service string, instances []string) error {
+	instancesBytes, err := json.Marshal(instances)
+	if err != nil {
+		return err
+	}
+
+	err = badgerDB.Update(func(txn *badger.Txn) error {
+		err := txn.Set([]byte(service), instancesBytes)
+		return err
+	})
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func GetInstance(service string) (string, error) {
+	instances, err := GetAllInstances(service)
+	if err != nil || len(instances) == 0 {
+		if err != nil {
+			log.Println("DB err:", err)
+		}
+
+		link, ok := services.FallbackMap[service]
+		if !ok {
+			return "", errors.New("invalid service")
+		}
+
+		return link, nil
+	}
+
+	previous, ok := selectionMap[service]
+	if ok && len(instances) > 2 {
+		instances = slices.DeleteFunc(instances, func(i string) bool {
+			return i == previous
+		})
+	}
+
+	index := rand.Intn(len(instances))
+	value := instances[index]
+	selectionMap[service] = value
+	return value, nil
+}
+
+func GetAllInstances(service string) ([]string, error) {
+	var instances []string
+	err := badgerDB.View(func(txn *badger.Txn) error {
+		item, err := txn.Get([]byte(service))
+		if err != nil {
+			return err
+		}
+
+		err = item.Value(func(val []byte) error {
+			err := json.Unmarshal(val, &instances)
+			return err
+		})
+
+		return err
+	})
+
+	return instances, err
+}
+
+func GetServiceList() []services.Service {
+	if cacheUpdated.Add(5 * time.Minute).After(time.Now().UTC()) {
+		return cachedServiceList
+	}
+
+	canCache := true
+
+	var serviceList []services.Service
+	for _, service := range services.ServiceList {
+		instances, err := GetAllInstances(service.Type)
+		if err != nil {
+			canCache = false
+			instances = []string{service.Fallback}
+		}
+
+		storedService := services.Service{
+			Type:      service.Type,
+			Instances: instances,
+		}
+
+		serviceList = append(serviceList, storedService)
+	}
+
+	if canCache {
+		cachedServiceList = serviceList
+		cacheUpdated = time.Now().UTC()
+	}
+
+	return serviceList
+}
+
+func CloseDB() error {
+	log.Println("Closing database...")
+	err := badgerDB.Close()
+	if err != nil {
+		log.Println("Error closing database", err)
+		return err
+	}
+
+	log.Println("Database closed!")
+	return nil
+}
+
+func init() {
+	selectionMap = make(map[string]string)
+}
diff --git a/db/db_test.go b/db/db_test.go
new file mode 100644
index 0000000..7bc92ca
--- /dev/null
+++ b/db/db_test.go
@@ -0,0 +1,60 @@
+package db
+
+import (
+	"log"
+	"os"
+	"slices"
+	"testing"
+)
+
+func TestMain(m *testing.M) {
+	err := InitializeDB()
+	if err != nil {
+		log.Fatalln("Failed to initialize database")
+	}
+
+	exitCode := m.Run()
+
+	_ = CloseDB()
+	os.Exit(exitCode)
+}
+
+func TestDatabase(t *testing.T) {
+	var (
+		service = "test"
+		siteA   = "a.com"
+		siteB   = "b.com"
+		siteC   = "c.com"
+	)
+	instances := []string{siteA, siteB, siteC}
+	err := SetInstances(service, instances)
+	if err != nil {
+		t.Fatalf("Failed to set instances: %v\n", err)
+	}
+
+	dbInstances, err := GetAllInstances(service)
+	if err != nil {
+		t.Fatalf("Failed to retrieve instances: %v\n", err)
+	}
+
+	for _, instance := range instances {
+		idx := slices.Index(dbInstances, instance)
+		if idx < 0 {
+			t.Fatalf("Failed to find instance in list")
+		}
+	}
+
+	firstInstance, err := GetInstance(service)
+	if err != nil {
+		t.Fatalf("Failed to fetch single instance: %v\n", err)
+	}
+
+	secondInstance, err := GetInstance(service)
+	if err != nil {
+		t.Fatalf("Failed to fetch single instance (second): %v\n", err)
+	} else if firstInstance == secondInstance {
+		t.Fatalf("Same instance was selected twice")
+	}
+
+	_ = CloseDB()
+}