Queued request timeout (prebid#1217)

Co-authored-by: Veronika Solovei <veronika.solovei@xandr.com>

config/config.go

@@ -64,6 +64,8 @@ type Configuration struct {
 	AccountRequired bool `mapstructure:"account_required"`
 	// Local private file containing SSL certificates
 	PemCertsFile string `mapstructure:"certificates_file"`
+	// Custom headers to handle request timeouts from queueing infrastructure
+	RequestTimeoutHeaders RequestTimeoutHeaders `mapstructure:"request_timeout_headers"`
 }
 
 const MIN_COOKIE_SIZE_BYTES = 500
@@ -199,6 +201,11 @@ type HostCookie struct {
 	TTL int64 `mapstructure:"ttl_days"`
 }
 
+type RequestTimeoutHeaders struct {
+	RequestTimeInQueue    string `mapstructure:"request_time_in_queue"`
+	RequestTimeoutInQueue string `mapstructure:"request_timeout_in_queue"`
+}
+
 func (cfg *HostCookie) TTLDuration() time.Duration {
 	return time.Duration(cfg.TTL) * time.Hour * 24
 }
@@ -748,6 +755,9 @@ func SetupViper(v *viper.Viper, filename string) {
 	v.SetDefault("account_required", false)
 	v.SetDefault("certificates_file", "")
 
+	v.SetDefault("request_timeout_headers.request_time_in_queue", "")
+	v.SetDefault("request_timeout_headers.request_timeout_in_queue", "")
+
 	// Set environment variable support:
 	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
 	v.SetEnvPrefix("PBS")

router/aspects/request_timeout_handler.go

@@ -0,0 +1,43 @@
+package aspects
+
+import (
+	"github.com/julienschmidt/httprouter"
+	"github.com/prebid/prebid-server/config"
+	"net/http"
+	"strconv"
+)
+
+func QueuedRequestTimeout(f httprouter.Handle, reqTimeoutHeaders config.RequestTimeoutHeaders) httprouter.Handle {
+
+	return func(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
+
+		reqTimeInQueue := r.Header.Get(reqTimeoutHeaders.RequestTimeInQueue)
+		reqTimeout := r.Header.Get(reqTimeoutHeaders.RequestTimeoutInQueue)
+
+		//If request timeout headers are not specified - process request as usual
+		if reqTimeInQueue == "" || reqTimeout == "" {
+			f(w, r, params)
+			return
+		}
+
+		reqTimeFloat, reqTimeFloatErr := strconv.ParseFloat(reqTimeInQueue, 64)
+		reqTimeoutFloat, reqTimeoutFloatErr := strconv.ParseFloat(reqTimeout, 64)
+
+		//Return HTTP 500 if request timeout headers are incorrect (wrong format)
+		if reqTimeFloatErr != nil || reqTimeoutFloatErr != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte("Request timeout headers are incorrect (wrong format)"))
+			return
+		}
+
+		//Return HTTP 408 if requests stays too long in queue
+		if reqTimeFloat >= reqTimeoutFloat {
+			w.WriteHeader(http.StatusRequestTimeout)
+			w.Write([]byte("Queued request processing time exceeded maximum"))
+			return
+		}
+
+		f(w, r, params)
+	}
+
+}

router/aspects/request_timeout_handler_test.go

@@ -0,0 +1,124 @@
+package aspects
+
+import (
+	"github.com/julienschmidt/httprouter"
+	"github.com/prebid/prebid-server/config"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const reqTimeInQueueHeaderName = "X-Ngx-Request-Time"
+const reqTimeoutHeaderName = "X-Request-Timeout"
+
+func TestAny(t *testing.T) {
+	testCases := []struct {
+		reqTimeInQueue          string
+		reqTimeOut              string
+		setHeaders              bool
+		extectedRespCode        int
+		expectedRespCodeMessage string
+		expectedRespBody        string
+		expectedRespBodyMessage string
+	}{
+		{
+			//TestQueuedRequestTimeoutWithTimeout
+			reqTimeInQueue:          "6",
+			reqTimeOut:              "5",
+			setHeaders:              true,
+			extectedRespCode:        http.StatusRequestTimeout,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 408",
+			expectedRespBody:        "Queued request processing time exceeded maximum",
+			expectedRespBodyMessage: "Body should have error message",
+		},
+		{
+			//TestQueuedRequestTimeoutNoTimeout
+			reqTimeInQueue:          "0.9",
+			reqTimeOut:              "5",
+			setHeaders:              true,
+			extectedRespCode:        http.StatusOK,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 200",
+			expectedRespBody:        "Executed",
+			expectedRespBodyMessage: "Body should be present in response",
+		},
+		{
+			//TestQueuedRequestNoHeaders
+			reqTimeInQueue:          "",
+			reqTimeOut:              "",
+			setHeaders:              false,
+			extectedRespCode:        http.StatusOK,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 200",
+			expectedRespBody:        "Executed",
+			expectedRespBodyMessage: "Body should be present in response",
+		},
+		{
+			//TestQueuedRequestSomeHeaders
+			reqTimeInQueue:          "2",
+			reqTimeOut:              "",
+			setHeaders:              true,
+			extectedRespCode:        http.StatusOK,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 200",
+			expectedRespBody:        "Executed",
+			expectedRespBodyMessage: "Body should be present in response",
+		},
+		{
+			//TestQueuedRequestAllHeadersIncorrect
+			reqTimeInQueue:          "test1",
+			reqTimeOut:              "test2",
+			setHeaders:              true,
+			extectedRespCode:        http.StatusInternalServerError,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 400",
+			expectedRespBody:        "Request timeout headers are incorrect (wrong format)",
+			expectedRespBodyMessage: "Body should have error message",
+		},
+		{
+			//TestQueuedRequestSomeHeadersIncorrect
+			reqTimeInQueue:          "test1",
+			reqTimeOut:              "123",
+			setHeaders:              true,
+			extectedRespCode:        http.StatusInternalServerError,
+			expectedRespCodeMessage: "Http response code is incorrect, should be 400",
+			expectedRespBody:        "Request timeout headers are incorrect (wrong format)",
+			expectedRespBodyMessage: "Body should have error message",
+		},
+	}
+
+	for _, test := range testCases {
+		result := ExecuteAspectRequest(t, test.reqTimeInQueue, test.reqTimeOut, test.setHeaders)
+		assert.Equal(t, test.extectedRespCode, result.Code, test.expectedRespCodeMessage)
+		assert.Equal(t, test.expectedRespBody, string(result.Body.Bytes()), test.expectedRespBodyMessage)
+	}
+}
+
+func MockEndpoint() httprouter.Handle {
+	return httprouter.Handle(MockHandler)
+}
+
+func MockHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+	w.Write([]byte("Executed"))
+}
+
+func ExecuteAspectRequest(t *testing.T, timeInQueue string, reqTimeout string, setHeaders bool) *httptest.ResponseRecorder {
+	rw := httptest.NewRecorder()
+	req, err := http.NewRequest("POST", "/test", nil)
+	if err != nil {
+		assert.Fail(t, "Unable create mock http request")
+	}
+	if setHeaders {
+		req.Header.Set(reqTimeInQueueHeaderName, timeInQueue)
+		req.Header.Set(reqTimeoutHeaderName, reqTimeout)
+	}
+
+	customHeaders := config.RequestTimeoutHeaders{reqTimeInQueueHeaderName, reqTimeoutHeaderName}
+
+	handler := QueuedRequestTimeout(MockEndpoint(), customHeaders)
+
+	r := httprouter.New()
+	r.POST("/test", handler)
+
+	r.ServeHTTP(rw, req)
+
+	return rw
+}

router/router.go

@@ -38,6 +38,7 @@ import (
 	"github.com/prebid/prebid-server/pbs"
 	metricsConf "github.com/prebid/prebid-server/pbsmetrics/config"
 	pbc "github.com/prebid/prebid-server/prebid_cache_client"
+	"github.com/prebid/prebid-server/router/aspects"
 	"github.com/prebid/prebid-server/ssl"
 	storedRequestsConf "github.com/prebid/prebid-server/stored_requests/config"
 	"github.com/prebid/prebid-server/usersync/usersyncers"
@@ -255,6 +256,11 @@ func New(cfg *config.Configuration, rateConvertor *currencies.RateConverter) (r
 		glog.Fatalf("Failed to create the video endpoint handler. %v", err)
 	}
 
+	requestTimeoutHeaders := config.RequestTimeoutHeaders{}
+	if cfg.RequestTimeoutHeaders != requestTimeoutHeaders {
+		videoEndpoint = aspects.QueuedRequestTimeout(videoEndpoint, cfg.RequestTimeoutHeaders)
+	}
+
 	r.POST("/auction", endpoints.Auction(cfg, syncers, gdprPerms, r.MetricsEngine, dataCache, exchanges))
 	r.POST("/openrtb2/auction", openrtbEndpoint)
 	r.POST("/openrtb2/video", videoEndpoint)