feat(callback): RetryAfter

Author: SevereCloud

callback/callback.go

@@ -12,6 +12,7 @@ import (
 	"log"
 	"net/http"
 	"strconv"
+	"time"
 
 	"github.com/SevereCloud/vksdk/v2/events"
 	"github.com/SevereCloud/vksdk/v2/internal"
@@ -81,13 +82,31 @@ func (cb *Callback) HandleFunc(w http.ResponseWriter, r *http.Request) {
 	retryCounter, _ := strconv.Atoi(r.Header.Get("X-Retry-Counter"))
 	ctx = context.WithValue(ctx, internal.CallbackRetryCounterKey, retryCounter)
 
+	var (
+		code int
+		date time.Time
+	)
+
+	retryAfter := func(c int, d time.Time) {
+		code = c
+		date = d
+	}
+	ctx = context.WithValue(ctx, internal.CallbackRetryAfterKey, retryAfter)
+
 	if err := cb.Handler(ctx, e); err != nil {
 		cb.logf("callback: %v", err)
 		http.Error(w, "Bad Request", http.StatusBadRequest)
 
 		return
 	}
 
+	if code != 0 {
+		w.Header().Set("Retry-After", date.Format(http.TimeFormat)) // RFC 7231, 7.1.3
+		http.Error(w, http.StatusText(code), code)
+
+		return
+	}
+
 	_, _ = w.Write([]byte("ok"))
 }
 

callback/context.go

@@ -2,6 +2,7 @@ package callback // import "github.com/SevereCloud/vksdk/v2/callback"
 
 import (
 	"context"
+	"time"
 
 	"github.com/SevereCloud/vksdk/v2/internal"
 )
@@ -10,3 +11,22 @@ import (
 func RetryCounterFromContext(ctx context.Context) int {
 	return ctx.Value(internal.CallbackRetryCounterKey).(int)
 }
+
+// RetryAfter send the "Retry-After" header field to indicate how long the
+// VK Callback ought to wait before making a repeated request.
+// Not work with goroutine mode!
+//
+// Possible HTTP status codes:
+//
+// 	http.StatusGone
+// 	http.StatusTooManyRequests
+// 	http.StatusServiceUnavailable
+//
+// The resend time range must be less than 3 hours. The actual time of
+// forwarding an event notification may be longer than the specified time.
+func RetryAfter(ctx context.Context, code int, date time.Time) {
+	ctx.Value(internal.CallbackRetryAfterKey).(func(int, time.Time))(
+		code,
+		date,
+	)
+}

callback/context_test.go

@@ -1,10 +1,15 @@
 package callback_test
 
 import (
+	"bytes"
 	"context"
+	"net/http"
+	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/SevereCloud/vksdk/v2/callback"
+	"github.com/SevereCloud/vksdk/v2/events"
 	"github.com/SevereCloud/vksdk/v2/internal"
 	"github.com/stretchr/testify/assert"
 )
@@ -18,3 +23,33 @@ func TestRetryCounterFromContext(t *testing.T) {
 	)
 	assert.Equal(t, retryCounter, callback.RetryCounterFromContext(ctx))
 }
+
+func TestRetryAfter(t *testing.T) {
+	code := http.StatusGone
+	date := time.Now().Add(time.Minute * 5)
+
+	cb := callback.NewCallback()
+	cb.MessageNew(func(ctx context.Context, obj events.MessageNewObject) {
+		callback.RetryAfter(
+			ctx,
+			code,
+			date,
+		)
+	})
+
+	jsonStr := []byte(`{"type": "message_new","object": {}}`)
+
+	req, err := http.NewRequest("POST", "/callback", bytes.NewBuffer(jsonStr))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rr := httptest.NewRecorder()
+	handler := http.HandlerFunc(cb.HandleFunc)
+
+	handler.ServeHTTP(rr, req)
+
+	assert.Equal(t, http.StatusText(code)+"\n", rr.Body.String())
+	assert.Equal(t, code, rr.Code)
+	assert.Equal(t, date.Format(http.TimeFormat), rr.Header().Get("Retry-After"))
+}

internal/transport.go

@@ -26,6 +26,7 @@ const (
 	EventIDKey
 	LongPollTsKey
 	CallbackRetryCounterKey
+	CallbackRetryAfterKey
 )
 
 // ContextClient return *http.Client.