From 3dfee77c4e5588c1c4fa4c7871d424310e9053db Mon Sep 17 00:00:00 2001
From: Anshal Shukla <shukla.anshal85@gmail.com>
Date: Sat, 11 Feb 2023 14:48:11 +0530
Subject: [PATCH] Modify queue-proxy to support Firecracker microVMs

---
 pkg/queue/sharedmain/main.go | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/pkg/queue/sharedmain/main.go b/pkg/queue/sharedmain/main.go
index 9b3c8e03005c..9e0937b85c48 100644
--- a/pkg/queue/sharedmain/main.go
+++ b/pkg/queue/sharedmain/main.go
@@ -31,7 +31,9 @@ import (
 	"go.uber.org/automaxprocs/maxprocs"
 	"go.uber.org/zap"
 
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/intstr"
 
 	"knative.dev/control-protocol/pkg/certificates"
 	netheader "knative.dev/networking/pkg/http/header"
@@ -114,6 +116,10 @@ type config struct {
 	ConcurrencyStateEndpoint  string `split_words:"true"` // optional
 	ConcurrencyStateTokenPath string `split_words:"true"` // optional
 
+	// vHive configuration
+	GuestAddr string `split_words:"true" required:"true"`
+	GuestPort string `split_words:"true" required:"true"`
+
 	Env
 }
 
@@ -233,6 +239,22 @@ func Main(opts ...Option) error {
 	// Setup probe to run for checking user-application healthiness.
 	// Do not set up probe if concurrency state endpoint is set, as
 	// paused containers don't play well with k8s readiness probes.
+	servingProbe := &corev1.Probe{
+		SuccessThreshold: 1,
+		ProbeHandler: corev1.ProbeHandler{
+			TCPSocket: &corev1.TCPSocketAction{
+				Host: env.GuestAddr,
+				Port: intstr.FromString(env.GuestPort),
+			},
+		},
+	}
+
+	var err error
+	env.ServingReadinessProbe, err = readiness.EncodeProbe(servingProbe)
+	if err != nil {
+		logger.Fatalw("Failed to create stats reporter", zap.Error(err))
+	}
+
 	probe := func() bool { return true }
 	if env.ServingReadinessProbe != "" && env.ConcurrencyStateEndpoint == "" {
 		probe = buildProbe(logger, env.ServingReadinessProbe, env.EnableHTTP2AutoDetection).ProbeContainer
@@ -343,7 +365,7 @@ func buildServer(ctx context.Context, env config, transport http.RoundTripper, p
 	ce *queue.ConcurrencyEndpoint, enableTLS bool) (*http.Server, *pkghandler.Drainer) {
 	// TODO: If TLS is enabled, execute probes twice and tracking two different sets of container health.
 
-	target := net.JoinHostPort("127.0.0.1", env.UserPort)
+	target := net.JoinHostPort(env.GuestAddr, env.GuestPort)
 
 	httpProxy := pkghttp.NewHeaderPruningReverseProxy(target, pkghttp.NoHostOverride, activator.RevisionHeaders, false /* use HTTP */)
 	httpProxy.Transport = transport