[YUNIKORN-2249] Add compression option to getQueueApplication API

pkg/webservice/handlers.go

@@ -19,6 +19,8 @@
 package webservice
 
 import (
+	"bytes"
+	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -59,6 +61,7 @@
 	GroupNameMissing         = "Group name is missing"
 	ApplicationDoesNotExists = "Application not found"
 	NodeDoesNotExists        = "Node not found"
+	UnsupportedCompType      = "Compression type not support"
 )
 
 var allowedActiveStatusMsg string
@@ -153,6 +156,10 @@
 	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)
@@ -455,7 +462,7 @@
 	}
 
 	if err := json.NewEncoder(w).Encode(result); err != nil {
 		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
 	}
 }
 
@@ -668,7 +675,7 @@
 	}
 	queueDao := queue.GetPartitionQueueDAOInfo(r.URL.Query().Has("subtree"))
 	if err := json.NewEncoder(w).Encode(queueDao); err != nil {
 		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
 	}
 }
 
@@ -745,9 +752,13 @@
 	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)
 	}
 }
 
@@ -1175,8 +1186,8 @@
 	if err := enc.Encode(dao.YunikornID{
 		InstanceUUID: schedulerContext.GetUUID(),
 	}); err != nil {
 		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
 	f.Flush()
 
@@ -1207,12 +1218,62 @@
 			}
 
 			if err := enc.Encode(e); err != nil {
 				log.Log(log.REST).Error("Marshalling error",
 					zap.String("host", r.Host))
 				buildJSONErrorResponse(w, err.Error(), http.StatusOK) // status code is 200 at this point, cannot be changed
 				return
 			}
 			f.Flush()
 		}
 	}
 }
+
+func checkHeader(h http.Header, key string, value string) bool {
+	values := h.Values(key)
+	for _, v := range values {
+		v2 := strings.Split(v, ",")
+		for _, item := range v2 {
+			item = strings.TrimSpace(item)
+			if item == 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
+	}
+
+	// don't compress the data if it is smaller than MTU size
+	if len(response) < 1500 {
+		if err = json.NewEncoder(w).Encode(data); err != nil {
+			buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		}
+		return
+	}
+
+	var compressedData bytes.Buffer
+	writer := gzip.NewWriter(&compressedData)
+	_, err = writer.Write(response)
+	if err != nil {
+		_ = writer.Close()
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	err = writer.Close()
+	if err != nil {
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	writeHeader(w, "Content-Encoding", "gzip")
+	if _, err = w.Write(compressedData.Bytes()); err != nil {
+		buildJSONErrorResponse(w, err.Error(), http.StatusInternalServerError)
+	}
+}

pkg/webservice/handlers_test.go

@@ -19,13 +19,17 @@
 package webservice
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -2600,3 +2604,59 @@ func NewResponseRecorderWithDeadline() *ResponseRecorderWithDeadline {
 		ResponseRecorder: httptest.NewRecorder(),
 	}
 }
+
+func TestCompressQueueApplications(t *testing.T) {
+	var appsDao []*dao.ApplicationDAOInfo
+
+	// case 1: data size is smaller than MTU size, so compression step is skipped
+	for i := 0; i < 2; i++ {
+		appName := "app-" + strconv.Itoa(i)
+		app := newApplication(appName, "part-01", "queue-1", rmID, security.UserGroup{})
+		appsDao = append(appsDao, getApplicationDAO(app))
+	}
+
+	resp := &MockResponseWriter{}
+	compress(resp, appsDao)
+
+	var decodedData []*dao.ApplicationDAOInfo
+	err := json.Unmarshal(resp.outputBytes, &decodedData)
+	assert.NilError(t, err, "Error when unmarshal data.")
+
+	for i := range decodedData {
+		assert.Equal(t, appsDao[i].ApplicationID, decodedData[i].ApplicationID)
+		assert.Equal(t, appsDao[i].Partition, decodedData[i].Partition)
+		assert.Equal(t, appsDao[i].QueueName, decodedData[i].QueueName)
+		assert.Equal(t, appsDao[i].SubmissionTime, decodedData[i].SubmissionTime)
+		assert.Equal(t, appsDao[i].User, decodedData[i].User)
+		assert.Equal(t, appsDao[i].Groups[0], decodedData[i].Groups[0])
+	}
+
+	// case 2: data size is greater than MTU size, so do the compression
+	for i := 2; i < 10; i++ {
+		appName := "app-" + strconv.Itoa(i)
+		app := newApplication(appName, "part-01", "queue-1", rmID, security.UserGroup{})
+		appsDao = append(appsDao, getApplicationDAO(app))
+	}
+	compress(resp, appsDao)
+
+	buf := bytes.NewBuffer(resp.outputBytes)
+	gzipReader, err := gzip.NewReader(buf)
+	assert.NilError(t, err, "Error while decompressing data.")
+	err = gzipReader.Close()
+	assert.NilError(t, err, "Error when close gzip reader.")
+
+	uncompressedData, err := io.ReadAll(gzipReader)
+	assert.NilError(t, err, "Error when read decoded data.")
+
+	err = json.Unmarshal(uncompressedData, &decodedData)
+	assert.NilError(t, err, "Error when unmarshal decoded data.")
+
+	for i := range decodedData {
+		assert.Equal(t, appsDao[i].ApplicationID, decodedData[i].ApplicationID)
+		assert.Equal(t, appsDao[i].Partition, decodedData[i].Partition)
+		assert.Equal(t, appsDao[i].QueueName, decodedData[i].QueueName)
+		assert.Equal(t, appsDao[i].SubmissionTime, decodedData[i].SubmissionTime)
+		assert.Equal(t, appsDao[i].User, decodedData[i].User)
+		assert.Equal(t, appsDao[i].Groups[0], decodedData[i].Groups[0])
+	}
+}