Implement WebSockets for Profile download (#2190)

restapi/admin_profiling.go

@@ -18,59 +18,32 @@ package restapi
 
 import (
 	"context"
-	"io"
+	"io/ioutil"
 	"net/http"
 
-	"github.com/go-openapi/runtime"
-	"github.com/go-openapi/runtime/middleware"
+	"github.com/gorilla/websocket"
 	"github.com/minio/console/models"
-	"github.com/minio/console/restapi/operations"
-	profileApi "github.com/minio/console/restapi/operations/profile"
 	"github.com/minio/madmin-go"
 )
 
-func registerProfilingHandler(api *operations.ConsoleAPI) {
-	// Start Profiling
-	api.ProfileProfilingStartHandler = profileApi.ProfilingStartHandlerFunc(func(params profileApi.ProfilingStartParams, session *models.Principal) middleware.Responder {
-		profilingStartResponse, err := getProfilingStartResponse(session, params)
-		if err != nil {
-			return profileApi.NewProfilingStartDefault(int(err.Code)).WithPayload(err)
-		}
-		return profileApi.NewProfilingStartCreated().WithPayload(profilingStartResponse)
-	})
-	// Stop and download profiling data
-	api.ProfileProfilingStopHandler = profileApi.ProfilingStopHandlerFunc(func(params profileApi.ProfilingStopParams, session *models.Principal) middleware.Responder {
-		profilingStopResponse, err := getProfilingStopResponse(session, params)
-		if err != nil {
-			return profileApi.NewProfilingStopDefault(int(err.Code)).WithPayload(err)
-		}
-		// Custom response writer to set the content-disposition header to tell the
-		// HTTP client the name and extension of the file we are returning
-		return middleware.ResponderFunc(func(w http.ResponseWriter, _ runtime.Producer) {
-			defer profilingStopResponse.Close()
-			w.Header().Set("Content-Type", "application/zip")
-			w.Header().Set("Content-Disposition", "attachment; filename=profile.zip")
-			io.Copy(w, profilingStopResponse)
-		})
-	})
+var items []*models.StartProfilingItem
+
+type profileOptions struct {
+	Types string
 }
 
-// startProfiling() starts the profiling on the Minio server
-// Enable 1 of the 7 profiling mechanisms: "cpu","mem","block","mutex","trace","threads","goroutines"
-// in the Minio server, returns []*models.StartProfilingItem that contains individual status of this operation
-// for each Minio node, ie:
-//
-//	{
-//		"Success": true,
-//		"nodeName": "127.0.0.1:9000"
-//		"errors": ""
-//	}
-func startProfiling(ctx context.Context, client MinioAdmin, profilerType string) ([]*models.StartProfilingItem, error) {
-	profilingResults, err := client.startProfiling(ctx, madmin.ProfilerType(profilerType))
+func getProfileOptionsFromReq(req *http.Request) (*profileOptions, error) {
+	pOptions := profileOptions{}
+	pOptions.Types = req.FormValue("types")
+	return &pOptions, nil
+}
+
+func startProfiling(ctx context.Context, conn WSConn, client MinioAdmin, pOpts *profileOptions) error {
+	profilingResults, err := client.startProfiling(ctx, madmin.ProfilerType(pOpts.Types))
 	if err != nil {
-		return nil, err
+		return err
 	}
-	var items []*models.StartProfilingItem
+	items = []*models.StartProfilingItem{}
 	for _, result := range profilingResults {
 		items = append(items, &models.StartProfilingItem{
 			Success:  result.Success,
@@ -78,57 +51,13 @@ func startProfiling(ctx context.Context, client MinioAdmin, profilerType string)
 			NodeName: result.NodeName,
 		})
 	}
-	return items, nil
-}
-
-// getProfilingStartResponse performs startProfiling() and serializes it to the handler's output
-func getProfilingStartResponse(session *models.Principal, params profileApi.ProfilingStartParams) (*models.StartProfilingList, *models.Error) {
-	ctx, cancel := context.WithCancel(params.HTTPRequest.Context())
-	defer cancel()
-	if params.Body == nil {
-		return nil, ErrorWithContext(ctx, ErrPolicyBodyNotInRequest)
-	}
-	mAdmin, err := NewMinioAdminClient(session)
-	if err != nil {
-		return nil, ErrorWithContext(ctx, err)
-	}
-	// create a MinIO Admin Client interface implementation
-	// defining the client to be used
-	adminClient := AdminClient{Client: mAdmin}
-	profilingItems, err := startProfiling(ctx, adminClient, *params.Body.Type)
-	if err != nil {
-		return nil, ErrorWithContext(ctx, err)
-	}
-	profilingList := &models.StartProfilingList{
-		StartResults: profilingItems,
-		Total:        int64(len(profilingItems)),
-	}
-	return profilingList, nil
-}
-
-// stopProfiling() stop the profiling on the Minio server and returns
-// the generated Zip file as io.ReadCloser
-func stopProfiling(ctx context.Context, client MinioAdmin) (io.ReadCloser, error) {
 	zippedData, err := client.stopProfiling(ctx)
 	if err != nil {
-		return nil, err
+		return err
 	}
-	return zippedData, nil
-}
-
-// getProfilingStopResponse() performs SetPolicy() and serializes it to the handler's output
-func getProfilingStopResponse(session *models.Principal, params profileApi.ProfilingStopParams) (io.ReadCloser, *models.Error) {
-	ctx := params.HTTPRequest.Context()
-	mAdmin, err := NewMinioAdminClient(session)
+	message, err := ioutil.ReadAll(zippedData)
 	if err != nil {
-		return nil, ErrorWithContext(ctx, err)
+		return err
 	}
-	// create a MinIO Admin Client interface implementation
-	// defining the client to be used
-	adminClient := AdminClient{Client: mAdmin}
-	profilingData, err := stopProfiling(ctx, adminClient)
-	if err != nil {
-		return nil, ErrorWithContext(ctx, err)
-	}
-	return profilingData, nil
+	return conn.writeMessage(websocket.BinaryMessage, message)
 }

restapi/admin_profiling_test.go

@@ -21,6 +21,8 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net/http"
+	"net/url"
 	"testing"
 
 	"github.com/minio/madmin-go"
@@ -32,22 +34,38 @@ var (
 	minioStopProfiling  func() (io.ReadCloser, error)
 )
 
-// mock function of startProfiling()
+// mock function for startProfiling()
 func (ac adminClientMock) startProfiling(ctx context.Context, profiler madmin.ProfilerType) ([]madmin.StartProfilingResult, error) {
 	return minioStartProfiling(profiler)
 }
 
-// mock function of stopProfiling()
+// mock function for stopProfiling()
 func (ac adminClientMock) stopProfiling(ctx context.Context) (io.ReadCloser, error) {
 	return minioStopProfiling()
 }
 
+// Implementing fake closingBuffer to mock stopProfiling() (io.ReadCloser, error)
+type ClosingBuffer struct {
+	*bytes.Buffer
+}
+
+// Implementing a fake Close function for io.ReadCloser
+func (cb *ClosingBuffer) Close() error {
+	return nil
+}
+
 func TestStartProfiling(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	assert := assert.New(t)
 	adminClient := adminClientMock{}
-	// Test-1 : startProfiling() Get response from Minio server with one profiling object
+	mockWSConn := mockConn{}
+	function := "startProfiling()"
+	testOptions := &profileOptions{
+		Types: "cpu",
+	}
+
+	// Test-1 : startProfiling() Get response from MinIO server with one profiling object without errors
 	// mock function response from startProfiling()
 	minioStartProfiling = func(profiler madmin.ProfilerType) ([]madmin.StartProfilingResult, error) {
 		return []madmin.StartProfilingResult{
@@ -63,56 +81,40 @@ func TestStartProfiling(t *testing.T) {
 			},
 		}, nil
 	}
-	function := "startProfiling()"
-	cpuProfiler := "cpu"
-	startProfilingResults, err := startProfiling(ctx, adminClient, cpuProfiler)
+	// mock function response from stopProfiling()
+	minioStopProfiling = func() (io.ReadCloser, error) {
+		return &ClosingBuffer{bytes.NewBufferString("In memory string eaeae")}, nil
+	}
+	// mock function response from mockConn.writeMessage()
+	connWriteMessageMock = func(messageType int, p []byte) error {
+		return nil
+	}
+	err := startProfiling(ctx, mockWSConn, adminClient, testOptions)
 	if err != nil {
 		t.Errorf("Failed on %s:, error occurred: %s", function, err.Error())
 	}
-	assert.Equal(2, len(startProfilingResults))
-	// Test-2 : startProfiling() Correctly handles errors returned by Minio
+	assert.Equal(err, nil)
+
+	// Test-2 : startProfiling() Correctly handles errors returned by MinIO
 	// mock function response from startProfiling()
 	minioStartProfiling = func(profiler madmin.ProfilerType) ([]madmin.StartProfilingResult, error) {
 		return nil, errors.New("error")
 	}
-	_, err = startProfiling(ctx, adminClient, cpuProfiler)
+	err = startProfiling(ctx, mockWSConn, adminClient, testOptions)
 	if assert.Error(err) {
 		assert.Equal("error", err.Error())
 	}
-}
 
-// Implementing fake closingBuffer need it to mock stopProfiling() (io.ReadCloser, error)
-type ClosingBuffer struct {
-	*bytes.Buffer
-}
-
-// Implementing a fake Close function for io.ReadCloser
-func (cb *ClosingBuffer) Close() error {
-	return nil
-}
-
-func TestStopProfiling(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	assert := assert.New(t)
-	adminClient := adminClientMock{}
-	// Test-1 : stopProfiling() Get response from Minio server and that response is a readCloser interface
-	// mock function response from startProfiling()
-	minioStopProfiling = func() (io.ReadCloser, error) {
-		return &ClosingBuffer{bytes.NewBufferString("In memory string eaeae")}, nil
+	// Test-3: getProfileOptionsFromReq() correctly returns profile options from request
+	u, _ := url.Parse("ws://localhost/ws/profile?types=cpu,mem,block,mutex,trace,threads,goroutines")
+	req := &http.Request{
+		URL: u,
 	}
-	function := "stopProfiling()"
-	_, err := stopProfiling(ctx, adminClient)
-	if err != nil {
-		t.Errorf("Failed on %s:, error occurred: %s", function, err.Error())
-	}
-	// Test-2 : stopProfiling() Correctly handles errors returned by Minio
-	// mock function response from stopProfiling()
-	minioStopProfiling = func() (io.ReadCloser, error) {
-		return nil, errors.New("error")
-	}
-	_, err = stopProfiling(ctx, adminClient)
-	if assert.Error(err) {
-		assert.Equal("error", err.Error())
+	opts, err := getProfileOptionsFromReq(req)
+	if assert.NoError(err) {
+		expectedOptions := profileOptions{
+			Types: "cpu,mem,block,mutex,trace,threads,goroutines",
+		}
+		assert.Equal(expectedOptions.Types, opts.Types)
 	}
 }

restapi/configure_console.go

@@ -118,8 +118,6 @@ func configureAPI(api *operations.ConsoleAPI) http.Handler {
 	registerBucketsLifecycleHandlers(api)
 	// Register service handlers
 	registerServiceHandlers(api)
-	// Register profiling handlers
-	registerProfilingHandler(api)
 	// Register session handlers
 	registerSessionHandlers(api)
 	// Register version handlers

restapi/ws_handle.go

@@ -250,7 +250,6 @@ func serveWS(w http.ResponseWriter, req *http.Request) {
 			closeWsConn(conn)
 			return
 		}
-
 		wsAdminClient, err := newWebSocketAdminClient(conn, session)
 		if err != nil {
 			ErrorWithContext(ctx, err)
@@ -258,6 +257,20 @@ func serveWS(w http.ResponseWriter, req *http.Request) {
 			return
 		}
 		go wsAdminClient.speedtest(ctx, speedtestOpts)
+	case strings.HasPrefix(wsPath, `/profile`):
+		pOptions, err := getProfileOptionsFromReq(req)
+		if err != nil {
+			ErrorWithContext(ctx, fmt.Errorf("error getting profile options: %v", err))
+			closeWsConn(conn)
+			return
+		}
+		wsAdminClient, err := newWebSocketAdminClient(conn, session)
+		if err != nil {
+			ErrorWithContext(ctx, err)
+			closeWsConn(conn)
+			return
+		}
+		go wsAdminClient.profile(ctx, pOptions)
 
 	default:
 		// path not found
@@ -464,6 +477,21 @@ func (wsc *wsAdminClient) speedtest(ctx context.Context, opts *madmin.SpeedtestO
 	sendWsCloseMessage(wsc.conn, err)
 }
 
+func (wsc *wsAdminClient) profile(ctx context.Context, opts *profileOptions) {
+	defer func() {
+		LogInfo("profile stopped")
+		// close connection after return
+		wsc.conn.close()
+	}()
+	LogInfo("profile started")
+
+	ctx = wsReadClientCtx(ctx, wsc.conn)
+
+	err := startProfiling(ctx, wsc.conn, wsc.client, opts)
+
+	sendWsCloseMessage(wsc.conn, err)
+}
+
 // sendWsCloseMessage sends Websocket Connection Close Message indicating the Status Code
 // see https://tools.ietf.org/html/rfc6455#page-45
 func sendWsCloseMessage(conn WSConn, err error) {