Skip to content

Commit

Permalink
Check 'Accept-Encoding' header
Browse files Browse the repository at this point in the history
  • Loading branch information
Ivaka committed Jan 26, 2024
1 parent 6a9320c commit dd964e1
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 14 deletions.
54 changes: 42 additions & 12 deletions collector/collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,11 @@ import (
"github.com/castai/egressd/pb"
)

var (
acceptEncoding = http.CanonicalHeaderKey("Accept-Encoding")
contentEncoding = http.CanonicalHeaderKey("Content-Encoding")
)

func CurrentTimeGetter() func() time.Time {
return func() time.Time {
return time.Now()
Expand Down Expand Up @@ -141,20 +146,19 @@ func (c *Collector) GetRawNetworkMetricsHandler(w http.ResponseWriter, req *http
return
}

writer, err := gzip.NewWriterLevel(w, gzip.BestCompression)
if err != nil {
c.log.Errorf("cannot create gzip writer: %v", err)
w.WriteHeader(http.StatusInternalServerError)
return
enc := req.Header.Get(acceptEncoding)
if strings.Contains(strings.ToLower(enc), "gzip") {
if err := c.writeGzipBody(w, batchBytes); err != nil {
c.log.Errorf("write batch %v", err)
return
}
} else {
if err := c.writePlainBody(w, batchBytes); err != nil {
c.log.Errorf("write batch %v", err)
return
}
}
defer writer.Close()

w.Header().Add("Content-Encoding", "gzip")

if _, err := writer.Write(batchBytes); err != nil {
c.log.Errorf("write batch: %v", err)
return
}
if c.cfg.SendTrafficDelta {
// reset metric tx/rx values, so only delta numbers will be sent with the next batch
for _, m := range c.podMetrics {
Expand All @@ -166,6 +170,32 @@ func (c *Collector) GetRawNetworkMetricsHandler(w http.ResponseWriter, req *http
}
}

func (c *Collector) writeGzipBody(w http.ResponseWriter, body []byte) error {
writer, err := gzip.NewWriterLevel(w, gzip.BestCompression)
if err != nil {
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return fmt.Errorf("cannot create gzip writer: %w", err)
}
defer writer.Close()

w.Header().Add(contentEncoding, "gzip")

if _, err := writer.Write(body); err != nil {
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return err
}

return nil
}

func (c *Collector) writePlainBody(w http.ResponseWriter, body []byte) error {
if _, err := w.Write(body); err != nil {
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return err
}
return nil
}

// collect aggregates conntract records into reduced pod metrics.
func (c *Collector) collect() error {
start := time.Now()
Expand Down
8 changes: 6 additions & 2 deletions collector/collector_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -422,7 +422,9 @@ func TestCollector__GetRawNetworkMetricsHandler(t *testing.T) {

w := httptest.NewRecorder()
// scrape pods metrics and prepare them to be sent
coll.GetRawNetworkMetricsHandler(w, &http.Request{})
coll.GetRawNetworkMetricsHandler(w, &http.Request{Header: http.Header{
"Accept-Encoding": []string{"gzip"},
}})
r.Equal(200, w.Code)

batch := &pb.RawNetworkMetricBatch{}
Expand Down Expand Up @@ -515,7 +517,9 @@ func TestCollector__GetRawNetworkMetricsHandler(t *testing.T) {

w := httptest.NewRecorder()
// scrape pods metrics and prepare them to be sent
coll.GetRawNetworkMetricsHandler(w, &http.Request{})
coll.GetRawNetworkMetricsHandler(w, &http.Request{Header: http.Header{
"Accept-Encoding": []string{"gzip"},
}})
r.Equal(200, w.Code)

batch := &pb.RawNetworkMetricBatch{}
Expand Down

0 comments on commit dd964e1

Please sign in to comment.