Add tests, refactor storage, optimize start up time for raft

internal/admin/admin_api.go

@@ -22,15 +22,16 @@ import (
 func StartAdminAPI(
 	version string, conf *config.DGateConfig,
 	logger *zap.Logger, cs changestate.ChangeState,
-) {
+) error {
 	if conf.AdminConfig == nil {
 		logger.Warn("Admin API is disabled")
-		return
+		return nil
 	}
 
 	mux := chi.NewRouter()
-	configureRoutes(mux, version,
-		logger.Named("routes"), cs, conf)
+	if err := configureRoutes(mux, version, logger.Named("routes"), cs, conf); err != nil {
+		return err
+	}
 
 	// Start HTTP Server
 	go func() {
@@ -44,6 +45,7 @@ func StartAdminAPI(
 			ErrorLog: zap.NewStdLog(adminHttpLogger),
 		}
 		if err := server.ListenAndServe(); err != nil {
+			logger.Error("Error starting admin api", zap.Error(err))
 			panic(err)
 		}
 	}()
@@ -143,4 +145,5 @@ func StartAdminAPI(
 			}()
 		}
 	}
+	return nil
 }

internal/admin/admin_fsm.go

@@ -11,18 +11,40 @@ import (
 	"go.uber.org/zap"
 )
 
-type dgateAdminFSM struct {
-	cs     changestate.ChangeState
-	logger *zap.Logger
+type AdminFSM struct {
+	cs      changestate.ChangeState
+	storage raft.StableStore
+	logger  *zap.Logger
+
+	localState *saveState
 }
 
-var _ raft.BatchingFSM = (*dgateAdminFSM)(nil)
+var _ raft.BatchingFSM = (*AdminFSM)(nil)
+
+type saveState struct {
+	AppliedIndex uint64 `json:"aindex"`
+}
 
-func newDGateAdminFSM(logger *zap.Logger, cs changestate.ChangeState) *dgateAdminFSM {
-	return &dgateAdminFSM{cs, logger}
+func newAdminFSM(
+	logger *zap.Logger,
+	storage raft.StableStore,
+	cs changestate.ChangeState,
+) raft.FSM {
+	fsm := &AdminFSM{cs, storage, logger, &saveState{}}
+	stateBytes, err := storage.Get([]byte("prev_state"))
+	if err != nil {
+		logger.Error("error getting prev_state", zap.Error(err))
+	} else if len(stateBytes) != 0 {
+		if err = json.Unmarshal(stateBytes, &fsm.localState); err != nil {
+			logger.Warn("corrupted state detected", zap.ByteString("prev_state", stateBytes))
+		} else {
+			logger.Info("found state in store", zap.Any("prev_state", fsm.localState))
+		}
+	}
+	return fsm
 }
 
-func (fsm *dgateAdminFSM) applyLog(log *raft.Log, reload bool) (*spec.ChangeLog, error) {
+func (fsm *AdminFSM) applyLog(log *raft.Log, reload bool) (*spec.ChangeLog, error) {
 	switch log.Type {
 	case raft.LogCommand:
 		var cl spec.ChangeLog
@@ -53,48 +75,77 @@ func (fsm *dgateAdminFSM) applyLog(log *raft.Log, reload bool) (*spec.ChangeLog,
 	return nil, nil
 }
 
-func (fsm *dgateAdminFSM) Apply(log *raft.Log) any {
-	rft := fsm.cs.Raft()
-	fsm.logger.Debug("apply single log",
-		zap.Uint64("applied", rft.AppliedIndex()),
-		zap.Uint64("commit", rft.CommitIndex()),
-		zap.Uint64("last", rft.LastIndex()),
-		zap.Uint64("logIndex", log.Index),
-	)
-	_, err := fsm.applyLog(log, true)
-	return err
+func (fsm *AdminFSM) Apply(log *raft.Log) any {
+	resps := fsm.ApplyBatch([]*raft.Log{log})
+	if len(resps) != 1 {
+		panic("apply batch not returning the correct number of responses")
+	}
+	return resps[0]
 }
 
-func (fsm *dgateAdminFSM) ApplyBatch(logs []*raft.Log) []any {
+func (fsm *AdminFSM) ApplyBatch(logs []*raft.Log) []any {
 	rft := fsm.cs.Raft()
-	lastIndex := len(logs) - 1
+	appliedIndex := rft.AppliedIndex()
+	lastLogIndex := logs[len(logs)-1].Index
 	fsm.logger.Debug("apply log batch",
-		zap.Uint64("applied", rft.AppliedIndex()),
+		zap.Uint64("applied", appliedIndex),
 		zap.Uint64("commit", rft.CommitIndex()),
 		zap.Uint64("last", rft.LastIndex()),
 		zap.Uint64("log[0]", logs[0].Index),
-		zap.Uint64("log[-1]", logs[lastIndex].Index),
+		zap.Uint64("log[-1]", lastLogIndex),
 		zap.Int("logs", len(logs)),
 	)
+
+	var err error
 	results := make([]any, len(logs))
+
 	for i, log := range logs {
-		// TODO: check to see if this can be optimized channels raft node provides
-		_, err := fsm.applyLog(log, lastIndex == i)
-		if err != nil {
+		isLast := len(logs)-1 == i
+		reload := fsm.shouldReload(log, isLast)
+		if _, err = fsm.applyLog(log, reload); err != nil {
 			fsm.logger.Error("Error applying log", zap.Error(err))
 			results[i] = err
 		}
 	}
+
+	if appliedIndex != 0 && lastLogIndex >= appliedIndex {
+		fsm.localState.AppliedIndex = lastLogIndex
+		if err = fsm.saveFSMState(); err != nil {
+			fsm.logger.Warn("failed to save applied index state",
+				zap.Uint64("applied_index", lastLogIndex),
+			)
+		}
+		fsm.cs.SetReady(true)
+	}
+
 	return results
 }
 
-func (fsm *dgateAdminFSM) Snapshot() (raft.FSMSnapshot, error) {
+func (fsm *AdminFSM) saveFSMState() error {
+	fsm.logger.Debug("saving localState",
+		zap.Any("data", fsm.localState),
+	)
+	stateBytes, err := json.Marshal(fsm.localState)
+	if err != nil {
+		return err
+	}
+	return fsm.storage.Set([]byte("prev_state"), stateBytes)
+}
+
+func (fsm *AdminFSM) shouldReload(log *raft.Log, reload bool) bool {
+	if reload {
+		return log.Index >= fsm.localState.AppliedIndex
+	}
+	return false
+}
+
+func (fsm *AdminFSM) Snapshot() (raft.FSMSnapshot, error) {
 	fsm.cs = nil
 	fsm.logger.Warn("snapshots not supported")
 	return nil, errors.New("snapshots not supported")
 }
 
-func (fsm *dgateAdminFSM) Restore(rc io.ReadCloser) error {
+func (fsm *AdminFSM) Restore(rc io.ReadCloser) error {
 	fsm.logger.Warn("snapshots not supported, cannot restore")
 	return nil
 }

internal/admin/admin_raft.go

@@ -101,7 +101,7 @@ func setupRaft(
 		adminConfig.Replication.AdvertScheme+"://(address)/raft",
 	)
 	fsmLogger := logger.Named("fsm")
-	adminFSM := newDGateAdminFSM(fsmLogger, cs)
+	adminFSM := newAdminFSM(fsmLogger, configStore, cs)
 	raftNode, err := raft.NewRaft(
 		raftConfig, adminFSM, logStore,
 		configStore, snapStore, transport,

internal/admin/admin_routes.go

@@ -1,6 +1,7 @@
 package admin
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
@@ -28,50 +29,51 @@ func configureRoutes(
 	logger *zap.Logger,
 	cs changestate.ChangeState,
 	conf *config.DGateConfig,
-) {
+) error {
 	adminConfig := conf.AdminConfig
-	server.Use(func(next http.Handler) http.Handler {
-		ipList := iplist.NewIPList()
-		for _, address := range adminConfig.AllowList {
-			if strings.Contains(address, "/") {
-				if err := ipList.AddCIDRString(address); err != nil {
-					panic(fmt.Sprintf("invalid cidr address in admin.allow_list: %s", address))
-				}
-			} else {
-				if err := ipList.AddIPString(address); err != nil {
-					panic(fmt.Sprintf("invalid ip address in admin.allow_list: %s", address))
-				}
+	ipList := iplist.NewIPList()
+	for _, address := range adminConfig.AllowList {
+		if strings.Contains(address, "/") {
+			if err := ipList.AddCIDRString(address); err != nil {
+				return fmt.Errorf("invalid cidr address in admin.allow_list: %s", address)
+			}
+		} else {
+			if err := ipList.AddIPString(address); err != nil {
+				return fmt.Errorf("invalid ip address in admin.allow_list: %s", address)
 			}
 		}
-		// basic auth
-		var userMap map[string]string
-		// key auth
-		var keyMap map[string]struct{}
+	}
+	// basic auth
+	var userMap map[string]string
+	// key auth
+	var keyMap map[string]struct{}
 
-		switch adminConfig.AuthMethod {
-		case config.AuthMethodBasicAuth:
-			userMap = make(map[string]string)
-			if len(adminConfig.BasicAuth.Users) > 0 {
-				for i, user := range adminConfig.BasicAuth.Users {
-					if user.Username == "" || user.Password == "" {
-						panic(fmt.Sprintf("both username and password are required: admin.basic_auth.users[%d]", i))
-					}
-					userMap[user.Username] = user.Password
+	switch adminConfig.AuthMethod {
+	case config.AuthMethodBasicAuth:
+		userMap = make(map[string]string)
+		if len(adminConfig.BasicAuth.Users) > 0 {
+			for i, user := range adminConfig.BasicAuth.Users {
+				if user.Username == "" || user.Password == "" {
+					return errors.New(fmt.Sprintf("both username and password are required: admin.basic_auth.users[%d]", i))
 				}
+				userMap[user.Username] = user.Password
 			}
-		case config.AuthMethodKeyAuth:
-			keyMap = make(map[string]struct{})
-			if adminConfig.KeyAuth != nil && len(adminConfig.KeyAuth.Keys) > 0 {
-				if adminConfig.KeyAuth.QueryParamName != "" && adminConfig.KeyAuth.HeaderName != "" {
-					panic("only one of admin.key_auth.query_param_name or admin.key_auth.header_name can be set")
-				}
-				for _, key := range adminConfig.KeyAuth.Keys {
-					keyMap[key] = struct{}{}
-				}
+		}
+	case config.AuthMethodKeyAuth:
+		keyMap = make(map[string]struct{})
+		if adminConfig.KeyAuth != nil && len(adminConfig.KeyAuth.Keys) > 0 {
+			if adminConfig.KeyAuth.QueryParamName != "" && adminConfig.KeyAuth.HeaderName != "" {
+				return errors.New("only one of admin.key_auth.query_param_name or admin.key_auth.header_name can be set")
+			}
+			for _, key := range adminConfig.KeyAuth.Keys {
+				keyMap[key] = struct{}{}
 			}
-		case config.AuthMethodJWTAuth:
-			panic("JWT Auth is not supported yet")
 		}
+	case config.AuthMethodJWTAuth:
+		return errors.New("JWT Auth is not supported yet")
+	}
+
+	server.Use(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if ipList.Len() > 0 {
 				remoteIp := util.GetTrustedIP(r,
@@ -178,6 +180,7 @@ func configureRoutes(
 			misc.Handle("/metrics", promhttp.Handler())
 		}
 	})
+	return nil
 }
 
 func setupMetricProvider(

internal/config/configtest/dgate_configs.go

@@ -168,3 +168,15 @@ func NewTestDGateConfig_DomainAndNamespaces2() *config.DGateConfig {
 	conf.DisableDefaultNamespace = false
 	return conf
 }
+
+func NewTestAdminConfig() *config.DGateConfig {
+	conf := NewTestDGateConfig()
+	conf.AdminConfig = &config.DGateAdminConfig{
+		Host: "localhost",
+		Port: 0,
+		TLS: &config.DGateTLSConfig{
+			Port: 0,
+		},
+	}
+	return conf
+}

internal/proxy/change_log.go

@@ -37,27 +37,24 @@ func (ps *ProxyState) processChangeLog(cl *spec.ChangeLog, reload, store bool) (
 						return
 					}
 				}
-				// in memory storage for state restarts
+				if len(ps.changeLogs) > 0 {
+					xcl := ps.changeLogs[len(ps.changeLogs)-1]
+					if xcl.ID == cl.ID {
+						if r := ps.Raft(); r != nil && r.State() == raft.Leader {
+							return
+						}
+						ps.logger.Error("duplicate change log",
+							zap.String("id", cl.ID),
+							zap.Stringer("cmd", cl.Cmd),
+						)
+						return
+					}
+				}
 				ps.changeLogs = append(ps.changeLogs, cl)
 			}
 		}()
 	}
 
-	if len(ps.changeLogs) > 0 {
-		xcl := ps.changeLogs[len(ps.changeLogs)-1]
-		if xcl.ID == cl.ID {
-			if r := ps.Raft(); r != nil && r.State() == raft.Leader {
-				// FYI: we still need to store the change log
-				return nil
-			}
-			ps.logger.Error("duplicate change log",
-				zap.String("id", cl.ID),
-				zap.Stringer("cmd", cl.Cmd),
-			)
-			return errors.New("duplicate change log")
-		}
-	}
-
 	// apply change log to the state
 	if !cl.Cmd.IsNoop() {
 		defer func() {
@@ -77,9 +74,7 @@ func (ps *ProxyState) processChangeLog(cl *spec.ChangeLog, reload, store bool) (
 				})
 			}
 		}()
-		if cl.Cmd.Resource() == spec.Documents && !store {
-			return nil
-		} else if err = ps.processResource(cl); err != nil {
+		if err = ps.processResource(cl); err != nil {
 			ps.logger.Error("decoding or processing change log", zap.Error(err))
 			return
 		}
@@ -334,11 +329,11 @@ func (ps *ProxyState) restoreFromChangeLogs(directApply bool) error {
 
 func (ps *ProxyState) compactChangeLogs(logs []*spec.ChangeLog) (int, error) {
 	removeList := compactChangeLogsRemoveList(ps.logger, sliceutil.SliceCopy(logs))
-	removed, err := ps.store.DeleteChangeLogs(removeList)
+	err := ps.store.DeleteChangeLogs(removeList)
 	if err != nil {
-		return removed, err
+		return 0, err
 	}
-	return removed, nil
+	return len(logs), nil
 }
 
 /*

internal/proxy/dynamic_proxy.go

@@ -227,7 +227,7 @@ func (ps *ProxyState) startProxyServer() {
 	cfg := ps.config.ProxyConfig
 	hostPort := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
 	ps.logger.Info("Starting proxy server on " + hostPort)
-	proxyHttpLogger := ps.Logger()
+	proxyHttpLogger := ps.logger.Named("http")
 	server := &http.Server{
 		Addr:     hostPort,
 		Handler:  ps,
@@ -324,7 +324,7 @@ func (ps *ProxyState) Stop() {
 
 	ps.logger.Info("Stopping proxy server")
 	defer os.Exit(0)
-	defer ps.Logger().Sync()
+	defer ps.logger.Sync()
 
 	ps.proxyLock.Lock()
 	defer ps.proxyLock.Unlock()