Merge branch '188-dle-restart-with-sessions' into 'master'

feat: do not lose clones when DLE is restarted (#188)

Closes #188

See merge request postgres-ai/database-lab!342

Author: agneum

cmd/database-lab/main.go

@@ -5,7 +5,6 @@
 // TODO(anatoly):
 // - Validate configs in all components.
 // - Tests.
-// - Don't kill clones on shutdown/start.
 
 package main
 
@@ -18,7 +17,6 @@ import (
 
 	"github.com/docker/docker/client"
 	"github.com/pkg/errors"
-	"github.com/rs/xid"
 
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/config"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/config/global"
@@ -35,25 +33,20 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/runners"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/srv"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/util/networks"
-	"gitlab.com/postgres-ai/database-lab/v2/version"
 )
 
 const (
 	shutdownTimeout = 30 * time.Second
 )
 
 func main() {
-	log.Msg("Database Lab version: ", version.GetVersion())
-
-	instanceID := xid.New().String()
-
-	log.Msg("Database Lab Instance ID:", instanceID)
-
-	cfg, err := config.LoadConfiguration(instanceID)
+	cfg, err := config.LoadConfiguration()
 	if err != nil {
 		log.Fatal(errors.WithMessage(err, "failed to parse config"))
 	}
 
+	log.Msg("Database Lab Instance ID:", cfg.Global.InstanceID)
+
 	runner := runners.NewLocalRunner(cfg.Provision.UseSudo)
 
 	pm := pool.NewPoolManager(&cfg.PoolManager, runner)
@@ -74,13 +67,13 @@ func main() {
 		log.Err("hostname is empty")
 	}
 
-	internalNetwork, err := networks.Setup(ctx, dockerCLI, instanceID, hostname)
+	internalNetworkID, err := networks.Setup(ctx, dockerCLI, cfg.Global.InstanceID, hostname)
 	if err != nil {
 		log.Errf(err.Error())
 		return
 	}
 
-	defer networks.Stop(dockerCLI, internalNetwork.ID)
+	defer networks.Stop(dockerCLI, internalNetworkID, hostname)
 
 	// Create a platform service to make requests to Platform.
 	platformSvc, err := platform.New(ctx, cfg.Platform)
@@ -89,6 +82,11 @@ func main() {
 		return
 	}
 
+	dbCfg := &resources.DB{
+		Username: cfg.Global.Database.User(),
+		DBName:   cfg.Global.Database.Name(),
+	}
+
 	emergencyShutdown := func() {
 		cancel()
 
@@ -104,40 +102,37 @@ func main() {
 	if err := retrievalSvc.Run(ctx); err != nil {
 		log.Err("Failed to run the data retrieval service:", err)
 		emergencyShutdown()
+
 		return
 	}
 
 	defer retrievalSvc.Stop()
 
-	dbCfg := &resources.DB{
-		Username: cfg.Global.Database.User(),
-		DBName:   cfg.Global.Database.Name(),
-	}
-
 	// Create a cloning service to provision new clones.
-	provisionSvc, err := provision.New(ctx, &cfg.Provision, dbCfg, dockerCLI, pm, internalNetwork.ID)
+	provisionSvc, err := provision.New(ctx, &cfg.Provision, dbCfg, dockerCLI, pm, internalNetworkID)
 	if err != nil {
 		log.Errf(errors.WithMessage(err, `error in the "provision" section of the config`).Error())
 	}
 
-	obsCh := make(chan string, 1)
+	observingChan := make(chan string, 1)
 
-	cloningSvc := cloning.NewBase(&cfg.Cloning, provisionSvc, obsCh)
+	cloningSvc := cloning.NewBase(&cfg.Cloning, provisionSvc, observingChan)
 	if err = cloningSvc.Run(ctx); err != nil {
 		log.Err(err)
 		emergencyShutdown()
+
 		return
 	}
 
 	obs := observer.NewObserver(dockerCLI, &cfg.Observer, platformSvc.Client, pm)
 	est := estimator.NewEstimator(&cfg.Estimator)
 
-	go removeObservingClones(obsCh, obs)
+	go removeObservingClones(observingChan, obs)
 
 	server := srv.NewServer(&cfg.Server, &cfg.Global, obs, cloningSvc, platformSvc, dockerCLI, est, pm)
 	shutdownCh := setShutdownListener()
 
-	go setReloadListener(ctx, instanceID, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, est, server)
+	go setReloadListener(ctx, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, est, server)
 
 	server.InitHandlers()
 
@@ -159,11 +154,12 @@ func main() {
 	}
 
 	shutdownDatabaseLabEngine(shutdownCtx, dockerCLI, cfg.Global, pm.Active().Pool())
+	cloningSvc.SaveClonesState()
 }
 
-func reloadConfig(ctx context.Context, instanceID string, provisionSvc *provision.Provisioner, retrievalSvc *retrieval.Retrieval,
+func reloadConfig(ctx context.Context, provisionSvc *provision.Provisioner, retrievalSvc *retrieval.Retrieval,
 	pm *pool.Manager, cloningSvc *cloning.Base, platformSvc *platform.Service, est *estimator.Estimator, server *srv.Server) error {
-	cfg, err := config.LoadConfiguration(instanceID)
+	cfg, err := config.LoadConfiguration()
 	if err != nil {
 		return err
 	}
@@ -200,15 +196,15 @@ func reloadConfig(ctx context.Context, instanceID string, provisionSvc *provisio
 	return nil
 }
 
-func setReloadListener(ctx context.Context, instanceID string, provisionSvc *provision.Provisioner, retrievalSvc *retrieval.Retrieval,
+func setReloadListener(ctx context.Context, provisionSvc *provision.Provisioner, retrievalSvc *retrieval.Retrieval,
 	pm *pool.Manager, cloningSvc *cloning.Base, platformSvc *platform.Service, est *estimator.Estimator, server *srv.Server) {
 	reloadCh := make(chan os.Signal, 1)
 	signal.Notify(reloadCh, syscall.SIGHUP)
 
 	for range reloadCh {
 		log.Msg("Reloading configuration")
 
-		if err := reloadConfig(ctx, instanceID, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, est, server); err != nil {
+		if err := reloadConfig(ctx, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, est, server); err != nil {
 			log.Err("Failed to reload configuration", err)
 		}
 

pkg/config/config.go

@@ -6,9 +6,12 @@
 package config
 
 import (
+	"fmt"
 	"os"
+	"path/filepath"
 
 	"github.com/pkg/errors"
+	"github.com/rs/xid"
 	"gopkg.in/yaml.v2"
 
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/config/global"
@@ -25,7 +28,8 @@ import (
 )
 
 const (
-	configName = "server.yml"
+	configName     = "server.yml"
+	instanceIDFile = "instance_id"
 )
 
 // Config contains a common database-lab configuration.
@@ -42,7 +46,7 @@ type Config struct {
 }
 
 // LoadConfiguration instances a new application configuration.
-func LoadConfiguration(instanceID string) (*Config, error) {
+func LoadConfiguration() (*Config, error) {
 	cfg, err := readConfig()
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to parse config")
@@ -51,9 +55,28 @@ func LoadConfiguration(instanceID string) (*Config, error) {
 	log.SetDebug(cfg.Global.Debug)
 	log.Dbg("Config loaded", cfg)
 
-	cfg.Global.InstanceID = instanceID
+	return cfg, cfg.loadInstanceID()
+}
+
+// loadInstanceID tries to make instance ID persistent across runs and load its value after restart
+func (cfg *Config) loadInstanceID() error {
+	idFilepath := filepath.Join(cfg.PoolManager.MountDir, instanceIDFile)
+
+	data, err := os.ReadFile(idFilepath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			cfg.Global.InstanceID = xid.New().String()
+			log.Dbg("no instance_id file was found, generate new instance ID", cfg.Global.InstanceID)
+
+			return os.WriteFile(idFilepath, []byte(cfg.Global.InstanceID), 0544)
+		}
+
+		return fmt.Errorf("failed to load instanceid, %w", err)
+	}
+
+	cfg.Global.InstanceID = string(data)
 
-	return cfg, nil
+	return nil
 }
 
 // readConfig reads application configuration.
@@ -73,5 +96,5 @@ func readConfig() (*Config, error) {
 		return nil, errors.WithMessagef(err, "error parsing %s config", configPath)
 	}
 
-	return cfg, nil
+	return cfg, cfg.loadInstanceID()
 }

pkg/config/config_test.go

@@ -0,0 +1,80 @@
+package config
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/rs/xid"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+)
+
+func TestLoadConfig(t *testing.T) {
+	suite.Run(t, &ConfigSuite{})
+}
+
+func copyFile(src, dst string, process func([]byte) []byte) error {
+	data, err := os.ReadFile(src)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(dst, process(data), 0600)
+}
+
+type ConfigSuite struct {
+	suite.Suite
+	oldCwd   string
+	mountDir string
+}
+
+func (s *ConfigSuite) SetupTest() {
+	t := s.T()
+
+	s.mountDir = t.TempDir()
+
+	t.Log(s.mountDir)
+
+	cwd, err := os.Getwd()
+	require.NoError(t, err)
+	s.oldCwd = cwd
+	require.NoError(t, os.Chdir(s.mountDir))
+
+	require.NoError(t, os.Mkdir("configs", 0700))
+	require.NoError(t, os.Mkdir("data", 0700))
+
+	exampleSrc := filepath.Join(cwd, "../../configs/config.example.logical_generic.yml")
+	testConfig := filepath.Join(s.mountDir, "configs/server.yml")
+
+	err = copyFile(exampleSrc, testConfig, func(data []byte) []byte {
+		return bytes.ReplaceAll(data, []byte("/var/lib/dblab"), []byte(s.mountDir))
+	})
+	require.NoError(t, err)
+}
+
+func (s *ConfigSuite) TearDownTest() {
+	s.Require().NoError(os.Chdir(s.oldCwd))
+}
+
+func (s *ConfigSuite) TestGenerateNewID() {
+	cfg, err := readConfig()
+	s.Require().NoError(err)
+	s.NotEmpty(cfg.Global.InstanceID)
+
+	data, err := os.ReadFile(filepath.Join(s.mountDir, "instance_id"))
+	s.NoError(err)
+	s.Equal(cfg.Global.InstanceID, string(data))
+}
+
+func (s *ConfigSuite) TestLoadInstanceID() {
+	instanceID := xid.New().String()
+
+	err := os.WriteFile(filepath.Join(s.mountDir, "instance_id"), []byte(instanceID), 0600)
+	s.Require().NoError(err)
+
+	cfg, err := readConfig()
+	s.NoError(err)
+	s.Equal(cfg.Global.InstanceID, instanceID)
+}

pkg/observer/sql.go

@@ -89,11 +89,10 @@ func buildConnectionString(clone *models.Clone, socketDir string) string {
 		db.DBName = defaults.DBName
 	}
 
-	return fmt.Sprintf(`host=%s port=%s user=%s password='%s' database='%s' application_name='%s'`,
+	return fmt.Sprintf(`host=%s port=%s user=%s database='%s' application_name='%s'`,
 		socketDir,
 		db.Port,
 		db.Username,
-		db.Password,
 		db.DBName,
 		observerApplicationName)
 }