[YUNIKORN-2249] Add compression option to getQueueApplication API

pkg/webservice/handlers.go

@@ -19,6 +19,7 @@
 package webservice
 
 import (
+	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -59,6 +60,7 @@ const (
 	GroupNameMissing         = "Group name is missing"
 	ApplicationDoesNotExists = "Application not found"
 	NodeDoesNotExists        = "Node not found"
+	UnsupportedCompType      = "Compression type not support"
 )
 
 var allowedActiveStatusMsg string
@@ -153,6 +155,10 @@ func writeHeaders(w http.ResponseWriter) {
 	w.Header().Set("Access-Control-Allow-Headers", "X-Requested-With,Content-Type,Accept,Origin")
 }
 
+func writeHeader(w http.ResponseWriter, key, val string) {
+	w.Header().Set(key, val)
+}
+
 func buildJSONErrorResponse(w http.ResponseWriter, detail string, code int) {
 	w.WriteHeader(code)
 	errorInfo := dao.NewYAPIError(nil, code, detail)
@@ -745,6 +751,10 @@ func getQueueApplications(w http.ResponseWriter, r *http.Request) {
 	for _, app := range queue.GetCopyOfApps() {
 		appsDao = append(appsDao, getApplicationDAO(app))
 	}
+	if checkHeader(r.Header, "Accept-Encoding", "gzip") {
+		compress(w, appsDao)
+		return
+	}
 
 	if err := json.NewEncoder(w).Encode(appsDao); err != nil {
 		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
@@ -1216,3 +1226,38 @@ func getStream(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 }
+
+func checkHeader(h http.Header, key string, value string) bool {
+	values := h.Values(key)
+	for _, v := range values {
+		if strings.Contains(v, value) {
+			return true
+		}
+	}
+	return false
+}
+
+func compress(w http.ResponseWriter, data any) {
+	response, err := json.Marshal(data)
+	if err != nil {
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	writer := gzip.NewWriter(w)
+	w.Header().Set("Content-Encoding", "gzip")
+	_, err = writer.Write(response)
+	if err != nil {
+		_ = writer.Close()
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	err = writer.Flush()
+	if err != nil {
+		_ = writer.Close()
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	_ = writer.Close()
+}

pkg/webservice/handlers_test.go

@@ -19,10 +19,13 @@
 package webservice
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
@@ -1469,6 +1472,27 @@ func TestGetQueueApplicationsHandler(t *testing.T) {
 	resp = &MockResponseWriter{}
 	getQueueApplications(resp, req)
 	assertParamsMissing(t, resp)
+
+	// test compression
+	var req4 *http.Request
+	req4, err = http.NewRequest("GET", "/ws/v1/partition/default/queue/root.default/applications", strings.NewReader(""))
+	req4 = req4.WithContext(context.WithValue(req.Context(), httprouter.ParamsKey, httprouter.Params{
+		httprouter.Param{Key: "partition", Value: partitionNameWithoutClusterID},
+		httprouter.Param{Key: "queue", Value: "root.default"},
+	}))
+	req4.Header.Set("Accept-Encoding", "gzip")
+	assert.NilError(t, err, "Get Queue Applications Handler request failed")
+	resp4 := &MockResponseWriter{}
+	getQueueApplications(resp4, req4)
+
+	var appsDao4 []*dao.ApplicationDAOInfo
+	buf := bytes.NewBuffer(resp4.outputBytes)
+	gzipReader, err := gzip.NewReader(buf)
+	assert.NilError(t, err, "Create gzip reader failed.")
+	uncompress, err := io.ReadAll(gzipReader)
+	assert.NilError(t, err, "Decode gzip data failed.")
+	err = json.Unmarshal(uncompress, &appsDao4)
+	assert.NilError(t, err, unmarshalError)
 }
 
 func checkLegalGetAppsRequest(t *testing.T, url string, params httprouter.Params, expected []*dao.ApplicationDAOInfo) {