Merge pull request #45 from mackee/callback_timeout

Add callback.timeout configuration.

README.md

@@ -26,6 +26,7 @@ callback:
   # When your application returning "HTTP/1.1 OK 200" in this callback, a connection upgrade to WebSocket.
   # However your returning other than "200" ("404", "403" and "500"...), a connection is disconnect.
   connect: "http://localhost:12346/connect"
+  timeout: 10s    # callback response timeout
 proxy_set_header:
   "X-Foo": "Foo"  # set to callback request header
   "X-Bar": ""     # remove from callback request header

config.go

@@ -41,8 +41,9 @@ type Config struct {
 }
 
 type Callback struct {
-	Connect string `yaml:"connect"`
-	Close   string `yaml:"close"`
+	Connect string        `yaml:"connect"`
+	Close   string        `yaml:"close"`
+	Timeout time.Duration `yaml:"timeout"`
 }
 
 func NewConfig(filename string) (*Config, error) {

config_test.go

@@ -11,6 +11,7 @@ session_header: "X-Kuiperbelt-Session-Key"
 port: 12345
 callback:
   connect: "http://localhost:12346/connect"
+  timeout: 5s
 send_timeout: 1s
 send_queue_size: 1
 `)
@@ -20,6 +21,7 @@ var TestConfig = Config{
 	Callback: Callback{
 		Connect: "http://localhost:12346/connect",
 		Close:   "",
+		Timeout: time.Second * 5,
 	},
 	SendTimeout:   time.Second,
 	SendQueueSize: 1,

server.go

@@ -154,6 +154,7 @@ func (s *WebSocketServer) ConnectCallbackHandler(w http.ResponseWriter, r *http.
 		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 		return nil, err
 	}
+
 	// copy all headers except "Connection", "Upgrade" and "Sec-Websocket*"
 	for n, values := range r.Header {
 		n := strings.ToLower(n)
@@ -175,6 +176,13 @@ func (s *WebSocketServer) ConnectCallbackHandler(w http.ResponseWriter, r *http.
 	}
 
 	callbackRequest.Header.Add(ENDPOINT_HEADER_NAME, s.Config.Endpoint)
+
+	// set callback timeout
+	if timeout := s.Config.Callback.Timeout; timeout != 0 {
+		ctx, cancel := context.WithTimeout(r.Context(), timeout)
+		defer cancel()
+		callbackRequest = callbackRequest.WithContext(ctx)
+	}
 	resp, err := callbackClient.Do(callbackRequest)
 	if err != nil {
 		http.Error(w, http.StatusText(http.StatusBadGateway), http.StatusBadGateway)

server_test.go

@@ -97,6 +97,16 @@ func (s *testSuccessConnectCallbackServer) CloseHandler(w http.ResponseWriter, r
 	io.WriteString(w, "")
 }
 
+func (s *testSuccessConnectCallbackServer) SlowHandler(w http.ResponseWriter, r *http.Request) {
+	s.mu.Lock()
+	s.isCallbacked = true
+	s.header = r.Header
+	s.mu.Unlock()
+	time.Sleep(10 * time.Second)
+	w.WriteHeader(http.StatusOK)
+	io.WriteString(w, "slow response")
+}
+
 func newTestWebSocketRequest(url string) (*http.Request, error) {
 	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
@@ -313,3 +323,32 @@ func TestWebSocketSession__IdleTimeout(t *testing.T) {
 		t.Error("connection timeout is not working")
 	}
 }
+
+func TestWebSocketServer__Handler__SlowCallback(t *testing.T) {
+	var pool SessionPool
+	callbackServer := new(testSuccessConnectCallbackServer)
+	tcc := httptest.NewServer(http.HandlerFunc(callbackServer.SlowHandler))
+
+	c := TestConfig
+	c.Callback.Connect = tcc.URL
+
+	server := NewWebSocketServer(c, NewStats(), &pool)
+
+	tc := httptest.NewServer(http.HandlerFunc(server.Handler))
+
+	req, err := newTestWebSocketRequest(tc.URL)
+	if err != nil {
+		t.Fatal("cannot create request error:", err)
+	}
+
+	client := new(http.Client)
+	resp, err := client.Do(req)
+	if err != nil {
+		t.Fatal("to server upgrade request unexpected error:", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusBadGateway {
+		t.Error("unexpected status code:", resp.StatusCode)
+	}
+}