Allow to specify a connection timeout

Author: mudler

core/cli/explorer.go

@@ -2,15 +2,17 @@ package cli
 
 import (
 	"context"
+	"time"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
 	"github.com/mudler/LocalAI/core/explorer"
 	"github.com/mudler/LocalAI/core/http"
 )
 
 type ExplorerCMD struct {
-	Address      string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"`
-	PoolDatabase string `env:"LOCALAI_POOL_DATABASE,POOL_DATABASE" default:"explorer.json" help:"Path to the pool database" group:"api"`
+	Address           string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"`
+	PoolDatabase      string `env:"LOCALAI_POOL_DATABASE,POOL_DATABASE" default:"explorer.json" help:"Path to the pool database" group:"api"`
+	ConnectionTimeout string `env:"LOCALAI_CONNECTION_TIMEOUT,CONNECTION_TIMEOUT" default:"2m" help:"Connection timeout for the explorer" group:"api"`
 }
 
 func (e *ExplorerCMD) Run(ctx *cliContext.Context) error {
@@ -20,7 +22,11 @@ func (e *ExplorerCMD) Run(ctx *cliContext.Context) error {
 		return err
 	}
 
-	ds := explorer.NewDiscoveryServer(db)
+	dur, err := time.ParseDuration(e.ConnectionTimeout)
+	if err != nil {
+		return err
+	}
+	ds := explorer.NewDiscoveryServer(db, dur)
 
 	go ds.Start(context.Background())
 	appHTTP := http.Explorer(db, ds)

core/explorer/discovery.go

@@ -13,8 +13,9 @@ import (
 
 type DiscoveryServer struct {
 	sync.Mutex
-	database     *Database
-	networkState *NetworkState
+	database       *Database
+	networkState   *NetworkState
+	connectionTime time.Duration
 }
 
 type NetworkState struct {
@@ -29,9 +30,13 @@ func (s *DiscoveryServer) NetworkState() *NetworkState {
 
 // NewDiscoveryServer creates a new DiscoveryServer with the given Database.
 // it keeps the db state in sync with the network state
-func NewDiscoveryServer(db *Database) *DiscoveryServer {
+func NewDiscoveryServer(db *Database, dur time.Duration) *DiscoveryServer {
+	if dur == 0 {
+		dur = 50 * time.Second
+	}
 	return &DiscoveryServer{
-		database: db,
+		database:       db,
+		connectionTime: dur,
 		networkState: &NetworkState{
 			Networks: map[string]Network{},
 		},
@@ -49,7 +54,7 @@ func (s *DiscoveryServer) runBackground() {
 	}
 
 	for _, token := range s.database.TokenList() {
-		c, cancel := context.WithTimeout(context.Background(), 50*time.Second)
+		c, cancel := context.WithTimeout(context.Background(), s.connectionTime)
 		defer cancel()
 
 		// Connect to the network