Adding a new flag to customize endpoint to send GCI commands

Author: danielfireman

README.md

@@ -12,7 +12,7 @@ GCI has two main parts: i) the GCI-Proxy -- a multiplatform, runtime-agnostic HT
 ## Running GCI Proxy
 
 ```bash
-./gci-proxy --port 3000 --url http://localhost:8080 --ygen=67108864 --tgen=6710886
+./gci-proxy --port 3000 --url http://localhost:8080 --ygen=67108864 --tgen=6710886 --endpoint="/gci"
 ```
 
 Where the flags:
@@ -21,6 +21,7 @@ Where the flags:
 * --tgen: size in bytes of the tenured generation, e.g. 67108864 (64MB)
 * --port: port which the gci-proxy will listen, e.g. 3000
 * --url: URL of the server being proxied, e.g. --url http://localhost:8080
+* --endpoint: (optional) custom HTTP endpoint to send GCI-specific commands, e.g. --endpoint=/__gci
 
 
 ## GCI Proxy Protocol

main.go

@@ -32,15 +32,16 @@ func main() {
 	// flags
 	port := flag.String("port", defaultPort, defaultPortUsage)
 	redirectURL := flag.String("url", defaultTarget, defaultTargetUsage)
-	yGen := flag.Uint64("ygen", 0, "Invalid young generation size.")
-	tGen := flag.Uint64("tgen", 0, "Invalid tenured generation size.")
+	cmdEndpoint := flag.String("endpoint", "/", "Custom endpoint to send GCI-specific commands. For example, '/gci'")
+	yGen := flag.Uint64("ygen", 0, "Young generation size, in bytes.")
+	tGen := flag.Uint64("tgen", 0, "Tenured generation size, in bytes.")
 	flag.Parse()
 
 	if *yGen == 0 || *tGen == 0 {
 		log.Fatalf("Neither ygen nor tgen can be 0. ygen:%d tgen:%d", *yGen, *tGen)
 	}
 
-	proxy := newProxy(*redirectURL, *yGen, *tGen)
+	proxy := newProxy(*redirectURL, *cmdEndpoint, *yGen, *tGen)
 	c := make(chan struct{}, 1)
 	router := httprouter.New()
 	router.HandlerFunc("GET", "/", func(w http.ResponseWriter, r *http.Request) {
@@ -72,6 +73,7 @@ type transport struct {
 	isAvailable int32
 	shed        uint64
 	target      string
+	cmdEndpoint string
 	waiter      pendingWaiter
 	window      sampleWindow
 	stGen1      sheddingThreshold
@@ -139,7 +141,7 @@ func (t *transport) checkHeap() {
 	arrived, finished := t.waiter.waitPending()
 	t.window.update(finished)
 
-	req, err := http.NewRequest("GET", fmt.Sprintf("%s/", t.target), nil)
+	req, err := http.NewRequest("GET", fmt.Sprintf("%s%s", t.target, t.cmdEndpoint), nil)
 	if err != nil {
 		panic(fmt.Sprintf("Err trying to build heap check request: %q\n", err))
 	}
@@ -189,7 +191,7 @@ func shouldGC(pending, usedBytes, st uint64) bool {
 }
 
 func (t *transport) gc(gen generation) {
-	req, err := http.NewRequest("GET", fmt.Sprintf("%s/", t.target), nil)
+	req, err := http.NewRequest("GET", fmt.Sprintf("%s%s", t.target, t.cmdEndpoint), nil)
 	if err != nil {
 		panic(fmt.Sprintf("Err trying to build gc request: %q\n", err))
 	}
@@ -222,19 +224,20 @@ func (p *proxy) handle(w http.ResponseWriter, r *http.Request) {
 	p.proxy.ServeHTTP(w, r)
 }
 
-func newProxy(target string, yGen, tGen uint64) *proxy {
+func newProxy(target, cmdEndpoint string, yGen, tGen uint64) *proxy {
 	url, _ := url.Parse(target)
 	p := httputil.NewSingleHostReverseProxy(url)
-	p.Transport = newTransport(target, yGen, tGen)
+	p.Transport = newTransport(target, cmdEndpoint, yGen, tGen)
 	return &proxy{target: url, proxy: p}
 }
 
-func newTransport(target string, yGen, tGen uint64) *transport {
+func newTransport(target, cmdEndpoint string, yGen, tGen uint64) *transport {
 	return &transport{
-		target: target,
-		window: newSampleWindow(),
-		stGen1: newSheddingThreshold(time.Now().UnixNano(), yGen),
-		stGen2: newSheddingThreshold(time.Now().UnixNano(), tGen),
+		target:      target,
+		cmdEndpoint: cmdEndpoint,
+		window:      newSampleWindow(),
+		stGen1:      newSheddingThreshold(time.Now().UnixNano(), yGen),
+		stGen2:      newSheddingThreshold(time.Now().UnixNano(), tGen),
 	}
 }
 

main_test.go

@@ -60,7 +60,7 @@ func TestProxyHandle(t *testing.T) {
 	}))
 	defer target.Close()
 
-	p := newProxy(target.URL, 1024, 1024)
+	p := newProxy(target.URL, "/", 1024, 1024)
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		p.handle(w, r)
@@ -85,24 +85,28 @@ func TestTransport_CheckHeapSize(t *testing.T) {
 	data := []struct {
 		msg      string
 		response string
+		endpoint string
 	}{
-		{"onlyGen1", "1"},
-		{"bothGens", fmt.Sprintf("1%s1", genSeparator)},
+		{"onlyGen1_defaultEndpoint", "1", "/"},
+		{"bothGens_defaultEndpoint", fmt.Sprintf("1%s1", genSeparator), "/"},
+		{"onlyGen1_customEndpoint", "1", "/gci"},
+		{"bothGens_customEndpoint", fmt.Sprintf("1%s1", genSeparator), "/gci"},
 	}
 	for _, d := range data {
 		t.Run(d.msg, func(t *testing.T) {
 			var wg sync.WaitGroup
 			gotGCIHeapCheck := false
 			target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				if r.Header.Get(gciHeader) == heapCheckHeader {
+				fmt.Println(r.URL.Path)
+				if r.Header.Get(gciHeader) == heapCheckHeader && r.URL.Path == d.endpoint {
 					fmt.Fprintf(w, d.response)
 					gotGCIHeapCheck = true
 				}
 				wg.Done()
 			}))
 			defer target.Close()
 
-			server := proxyServer(target.URL, 1024, 1024)
+			server := proxyServer(target.URL, d.endpoint, 1024, 1024)
 			defer server.Close()
 
 			fireReqs(t, &wg, server.URL)
@@ -129,7 +133,7 @@ func TestTransport_GC(t *testing.T) {
 			var wg sync.WaitGroup
 			gcRan := false
 			target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				if r.Header.Get(gciHeader) == heapCheckHeader {
+				if r.Header.Get(gciHeader) == heapCheckHeader && r.URL.Path == "/gci" {
 					fmt.Fprintf(w, d.response)
 				}
 				if r.Header.Get(gciHeader) == d.gen {
@@ -139,7 +143,7 @@ func TestTransport_GC(t *testing.T) {
 			}))
 			defer target.Close()
 
-			server := proxyServer(target.URL, 1024, 1024)
+			server := proxyServer(target.URL, "/gci", 1024, 1024)
 			defer server.Close()
 
 			wg.Add(1)
@@ -152,8 +156,8 @@ func TestTransport_GC(t *testing.T) {
 	}
 }
 
-func proxyServer(target string, gen1, gen2 uint64) *httptest.Server {
-	p := newProxy(target, gen1, gen2)
+func proxyServer(target, cmdEndpoint string, gen1, gen2 uint64) *httptest.Server {
+	p := newProxy(target, cmdEndpoint, gen1, gen2)
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		p.handle(w, r)
 	}))
@@ -178,7 +182,7 @@ func BenchmarkProxyHandle(b *testing.B) {
 	}))
 	defer target.Close()
 
-	p := newProxy(target.URL, 1024, 1024)
+	p := newProxy(target.URL, "/", 1024, 1024)
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		p.handle(w, r)