feat: enforce milvus dial timeout if set (vllm-project#503)

Signed-off-by: cryo <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>

Author: cryo-zd

src/semantic-router/pkg/cache/cache_test.go

@@ -1,6 +1,7 @@
 package cache_test
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
@@ -177,6 +178,44 @@ development:
 				})
 			})
 
+			Context("Milvus connection timeouts", func() {
+				It("should respect connection timeout when endpoint is unreachable", func() {
+					unreachableConfigPath := filepath.Join(tempDir, "milvus-unreachable.yaml")
+					unreachableHost := "10.255.255.1" // unroutable address to simulate a hanging dial
+					unreachableConfig := fmt.Sprintf(`
+connection:
+  host: "%s"
+  port: 19530
+  database: "test_cache"
+  timeout: 1
+`, unreachableHost)
+
+					err := os.WriteFile(unreachableConfigPath, []byte(unreachableConfig), 0o644)
+					Expect(err).NotTo(HaveOccurred())
+
+					done := make(chan struct{})
+					var cacheErr error
+
+					go func() {
+						defer GinkgoRecover()
+						_, cacheErr = cache.NewMilvusCache(cache.MilvusCacheOptions{
+							Enabled:             true,
+							SimilarityThreshold: 0.85,
+							TTLSeconds:          60,
+							ConfigPath:          unreachableConfigPath,
+						})
+						close(done)
+					}()
+
+					Eventually(done, 2*time.Second, 100*time.Millisecond).Should(BeClosed())
+					Expect(cacheErr).To(HaveOccurred())
+					Expect(cacheErr.Error()).To(Or(
+						ContainSubstring("context deadline exceeded"),
+						ContainSubstring("timeout"),
+					))
+				})
+			})
+
 			Context("with unsupported backend type", func() {
 				It("should return error for unsupported backend type", func() {
 					config := cache.CacheConfig{

src/semantic-router/pkg/cache/milvus_cache.go

@@ -138,7 +138,16 @@ func NewMilvusCache(options MilvusCacheOptions) (*MilvusCache, error) {
 	// Establish connection to Milvus server
 	connectionString := fmt.Sprintf("%s:%d", config.Connection.Host, config.Connection.Port)
 	observability.Debugf("MilvusCache: connecting to Milvus at %s", connectionString)
-	milvusClient, err := client.NewGrpcClient(context.Background(), connectionString)
+	dialCtx := context.Background()
+	var cancel context.CancelFunc
+	if config.Connection.Timeout > 0 {
+		// If a timeout is specified, apply it to the connection context
+		timeout := time.Duration(config.Connection.Timeout) * time.Second
+		dialCtx, cancel = context.WithTimeout(dialCtx, timeout)
+		defer cancel()
+		observability.Debugf("MilvusCache: connection timeout set to %s", timeout)
+	}
+	milvusClient, err := client.NewGrpcClient(dialCtx, connectionString)
 	if err != nil {
 		observability.Debugf("MilvusCache: failed to connect: %v", err)
 		return nil, fmt.Errorf("failed to create Milvus client: %w", err)