diff --git a/go/vt/tabletserver/query_list_test.go b/go/vt/tabletserver/query_list_test.go index d4a91fe6fc3..522dd975020 100644 --- a/go/vt/tabletserver/query_list_test.go +++ b/go/vt/tabletserver/query_list_test.go @@ -7,20 +7,25 @@ import ( ) type testConn struct { - id int64 - query string + id int64 + query string + killed bool } -func (tc testConn) Current() string { return tc.query } +func (tc *testConn) Current() string { return tc.query } -func (tc testConn) ID() int64 { return tc.id } +func (tc *testConn) ID() int64 { return tc.id } -func (tc testConn) Kill() {} +func (tc *testConn) Kill() { tc.killed = true } + +func (tc *testConn) IsKilled() bool { + return tc.killed +} func TestQueryList(t *testing.T) { ql := NewQueryList() connID := int64(1) - qd := NewQueryDetail(context.Background(), testConn{id: connID}) + qd := NewQueryDetail(context.Background(), &testConn{id: connID}) ql.Add(qd) if qd1, ok := ql.queryDetails[connID]; !ok || qd1.connID != connID { @@ -28,7 +33,7 @@ func TestQueryList(t *testing.T) { } conn2ID := int64(2) - qd2 := NewQueryDetail(context.Background(), testConn{id: conn2ID}) + qd2 := NewQueryDetail(context.Background(), &testConn{id: conn2ID}) ql.Add(qd2) rows := ql.GetQueryzRows() diff --git a/go/vt/tabletserver/queryctl.go b/go/vt/tabletserver/queryctl.go index a2f232c6577..e8b4eaee4b9 100644 --- a/go/vt/tabletserver/queryctl.go +++ b/go/vt/tabletserver/queryctl.go @@ -382,6 +382,15 @@ func (rqsc *realQueryServiceControl) registerDebugHealthHandler() { }) } +func (rqsc *realQueryServiceControl) registerStreamQueryzHandlers() { + http.HandleFunc("/streamqueryz", func(w http.ResponseWriter, r *http.Request) { + streamQueryzHandler(rqsc.sqlQueryRPCService.qe.streamQList, w, r) + }) + http.HandleFunc("/streamqueryz/terminate", func(w http.ResponseWriter, r *http.Request) { + streamQueryzTerminateHandler(rqsc.sqlQueryRPCService.qe.streamQList, w, r) + }) +} + func buildFmter(logger *streamlog.StreamLogger) func(url.Values, interface{}) string { type formatter interface { Format(url.Values) string diff --git a/go/vt/tabletserver/stream_queryz.go b/go/vt/tabletserver/stream_queryz.go index 70980cb66c3..c3096fce2d8 100644 --- a/go/vt/tabletserver/stream_queryz.go +++ b/go/vt/tabletserver/stream_queryz.go @@ -1,3 +1,7 @@ +// Copyright 2015, Google Inc. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + package tabletserver import ( @@ -35,58 +39,55 @@ var ( `)) ) -func (rqsc *realQueryServiceControl) registerStreamQueryzHandlers() { - streamqueryzHandler := func(w http.ResponseWriter, r *http.Request) { - if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil { - acl.SendError(w, err) - return - } - rows := rqsc.sqlQueryRPCService.qe.streamQList.GetQueryzRows() - if err := r.ParseForm(); err != nil { - http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError) - return - } - format := r.FormValue("format") - if format == "json" { - js, err := json.Marshal(rows) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - w.Header().Set("Content-Type", "application/json") - w.Write(js) - return - } - startHTMLTable(w) - defer endHTMLTable(w) - w.Write(streamqueryzHeader) - for i := range rows { - if err := streamqueryzTmpl.Execute(w, rows[i]); err != nil { - log.Errorf("streamlogz: couldn't execute template: %v", err) - } - } +func streamQueryzHandler(queryList *QueryList, w http.ResponseWriter, r *http.Request) { + if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil { + acl.SendError(w, err) + return } - - http.HandleFunc("/streamqueryz", streamqueryzHandler) - http.HandleFunc("/streamqueryz/terminate", func(w http.ResponseWriter, r *http.Request) { - if err := acl.CheckAccessHTTP(r, acl.ADMIN); err != nil { - acl.SendError(w, err) - return - } - if err := r.ParseForm(); err != nil { - http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError) - return - } - connID := r.FormValue("connID") - c, err := strconv.Atoi(connID) + rows := queryList.GetQueryzRows() + if err := r.ParseForm(); err != nil { + http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError) + return + } + format := r.FormValue("format") + if format == "json" { + js, err := json.Marshal(rows) if err != nil { - http.Error(w, "invalid connID", http.StatusInternalServerError) + http.Error(w, err.Error(), http.StatusInternalServerError) return } - if err = rqsc.sqlQueryRPCService.qe.streamQList.Terminate(int64(c)); err != nil { - http.Error(w, fmt.Sprintf("error: %v", err), http.StatusInternalServerError) - return + w.Header().Set("Content-Type", "application/json") + w.Write(js) + return + } + startHTMLTable(w) + defer endHTMLTable(w) + w.Write(streamqueryzHeader) + for i := range rows { + if err := streamqueryzTmpl.Execute(w, rows[i]); err != nil { + log.Errorf("streamlogz: couldn't execute template: %v", err) } - streamqueryzHandler(w, r) - }) + } +} + +func streamQueryzTerminateHandler(queryList *QueryList, w http.ResponseWriter, r *http.Request) { + if err := acl.CheckAccessHTTP(r, acl.ADMIN); err != nil { + acl.SendError(w, err) + return + } + if err := r.ParseForm(); err != nil { + http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError) + return + } + connID := r.FormValue("connID") + c, err := strconv.Atoi(connID) + if err != nil { + http.Error(w, "invalid connID", http.StatusInternalServerError) + return + } + if err = queryList.Terminate(int64(c)); err != nil { + http.Error(w, fmt.Sprintf("error: %v", err), http.StatusInternalServerError) + return + } + streamQueryzHandler(queryList, w, r) } diff --git a/go/vt/tabletserver/stream_queryz_test.go b/go/vt/tabletserver/stream_queryz_test.go new file mode 100644 index 00000000000..9f4a8a408bb --- /dev/null +++ b/go/vt/tabletserver/stream_queryz_test.go @@ -0,0 +1,95 @@ +// Copyright 2015, Google Inc. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package tabletserver + +import ( + "net/http" + "net/http/httptest" + "testing" + + "golang.org/x/net/context" +) + +func TestStreamQueryzHandlerJSON(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/streamqueryz?format=json", nil) + + queryList := NewQueryList() + queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 1})) + queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 2})) + + streamQueryzHandler(queryList, resp, req) +} + +func TestStreamQueryzHandlerHTTP(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/streamqueryz", nil) + + queryList := NewQueryList() + queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 1})) + queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 2})) + + streamQueryzHandler(queryList, resp, req) +} + +func TestStreamQueryzHandlerHTTPFailedInvalidForm(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/streamqueryz", nil) + + streamQueryzHandler(NewQueryList(), resp, req) + if resp.Code != http.StatusInternalServerError { + t.Fatalf("http call should fail and return code: %d, but got: %d", + http.StatusInternalServerError, resp.Code) + } +} + +func TestStreamQueryzHandlerTerminateConn(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=1", nil) + + queryList := NewQueryList() + testConn := &testConn{id: 1} + queryList.Add(NewQueryDetail(context.Background(), testConn)) + if testConn.IsKilled() { + t.Fatalf("conn should still be alive") + } + streamQueryzTerminateHandler(queryList, resp, req) + if !testConn.IsKilled() { + t.Fatalf("conn should be killed") + } +} + +func TestStreamQueryzHandlerTerminateFailedInvalidConnID(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=invalid", nil) + + streamQueryzTerminateHandler(NewQueryList(), resp, req) + if resp.Code != http.StatusInternalServerError { + t.Fatalf("http call should fail and return code: %d, but got: %d", + http.StatusInternalServerError, resp.Code) + } +} + +func TestStreamQueryzHandlerTerminateFailedKnownConnID(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=10", nil) + + streamQueryzTerminateHandler(NewQueryList(), resp, req) + if resp.Code != http.StatusInternalServerError { + t.Fatalf("http call should fail and return code: %d, but got: %d", + http.StatusInternalServerError, resp.Code) + } +} + +func TestStreamQueryzHandlerTerminateFailedInvalidForm(t *testing.T) { + resp := httptest.NewRecorder() + req, _ := http.NewRequest("POST", "/streamqueryz/terminate?inva+lid=2", nil) + + streamQueryzTerminateHandler(NewQueryList(), resp, req) + if resp.Code != http.StatusInternalServerError { + t.Fatalf("http call should fail and return code: %d, but got: %d", + http.StatusInternalServerError, resp.Code) + } +}