goc/pkg/server/api.go

/*
 Copyright 2021 Qiniu Cloud (qiniu.com)
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
     http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package server

import (
	"bytes"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/RickLeee/goc-v2/pkg/log"
	"github.com/gin-gonic/gin"
	"golang.org/x/tools/cover"
	"k8s.io/test-infra/gopherage/pkg/cov"
)

// listAgents return all service informations
func (gs *gocServer) listAgents(c *gin.Context) {
	idQuery := c.Query("id")
	ifInIdMap := idMaps(idQuery)

	agents := make([]*gocCoveredAgent, 0)

	gs.agents.Range(func(key, value interface{}) bool {
		// check if id is in the query ids
		if !ifInIdMap(key.(string)) {
			return true
		}

		agent, ok := value.(*gocCoveredAgent)
		if !ok {
			return false
		}
		agents = append(agents, agent)
		return true
	})

	c.JSON(http.StatusOK, gin.H{
		"items": agents,
	})
}

func (gs *gocServer) removeAgents(c *gin.Context) {
	idQuery := c.Query("id")
	ifInIdMap := idMaps(idQuery)

	errs := ""
	gs.agents.Range(func(key, value interface{}) bool {

		// check if id is in the query ids
		id := key.(string)
		if !ifInIdMap(id) {
			return true
		}

		agent, ok := value.(*gocCoveredAgent)
		if !ok {
			return false
		}

		err := gs.removeAgentFromStore(id)
		if err != nil {
			log.Errorf("fail to remove agent: %v", id)
			err := fmt.Errorf("fail to remove agent: %v, err: %v", id, err)
			errs = errs + err.Error()
			return true
		}
		agent.closeConnection()
		gs.agents.Delete(key)

		return true
	})

	if errs != "" {
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": errs,
		})
	} else {
		c.JSON(http.StatusOK, nil)
	}
}

// getProfiles get and merge all agents' informations
//
// it is synchronous
func (gs *gocServer) getProfiles(c *gin.Context) {
	idQuery := c.Query("id")
	ifInIdMap := idMaps(idQuery)

	skippatternRaw := c.Query("skippattern")
	var skippattern []string
	if skippatternRaw != "" {
		skippattern = strings.Split(skippatternRaw, ",")
	}
	neerpatternRaw := c.Query("needpattern")
	var neerpattern []string
	if neerpatternRaw != "" {
		neerpattern = strings.Split(neerpatternRaw, ",")
	}

	extra := c.Query("extra")
	isExtra := filterExtra(extra)

	var mu sync.Mutex
	var wg sync.WaitGroup

	mergedProfiles := make([][]*cover.Profile, 0)

	gs.agents.Range(func(key, value interface{}) bool {
		// check if id is in the query ids
		if !ifInIdMap(key.(string)) {
			// not in
			return true
		}

		agent, ok := value.(*gocCoveredAgent)
		if !ok {
			return false
		}

		// check if extra matches
		if !isExtra(agent.Extra) {
			// not match
			return true
		}

		wg.Add(1)
		// 并发 rpc，且每个 rpc 设超时时间 10 second
		go func() {
			defer wg.Done()

			timeout := time.Duration(10 * time.Second)
			done := make(chan error, 1)

			var req ProfileReq = "getprofile"
			var res ProfileRes
			go func() {
				// lock-free
				rpc := agent.rpc
				if rpc == nil || agent.Status == DISCONNECT {
					done <- nil
					return
				}
				err := agent.rpc.Call("GocAgent.GetProfile", req, &res)
				if err != nil {
					log.Errorf("fail to get profile from: %v, reasson: %v. let's close the connection", agent.Id, err)
				}
				done <- err
			}()

			select {
			// rpc 超时
			case <-time.After(timeout):
				log.Warnf("rpc call timeout: %v", agent.Hostname)
				// 关闭链接
				agent.closeRpcConnOnce()
			case err := <-done:
				// 调用 rpc 发生错误
				if err != nil {
					// 关闭链接
					agent.closeRpcConnOnce()
				}
			}
			// append profile
			profile, err := convertProfile([]byte(res))
			if err != nil {
				log.Errorf("fail to convert the received profile from: %v, reasson: %v. let's close the connection", agent.Id, err)
				// 关闭链接
				agent.closeRpcConnOnce()
				return
			}

			// check if skippattern matches
			newProfile := filterProfileByPattern(skippattern, neerpattern, profile)

			mu.Lock()
			defer mu.Unlock()
			mergedProfiles = append(mergedProfiles, newProfile)
		}()

		return true
	})

	// 一直等待并发的 rpc 都回应
	wg.Wait()

	merged, err := cov.MergeMultipleProfiles(mergedProfiles)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": err.Error(),
		})
		return
	}

	var buff bytes.Buffer
	err = cov.DumpProfile(merged, &buff)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": err.Error(),
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"profile": buff.String(),
	})
}

// resetProfiles reset all profiles in agent
//
// it is async, the function will return immediately
func (gs *gocServer) resetProfiles(c *gin.Context) {
	idQuery := c.Query("id")
	ifInIdMap := idMaps(idQuery)

	extra := c.Query("extra")
	isExtra := filterExtra(extra)

	gs.agents.Range(func(key, value interface{}) bool {

		// check if id is in the query ids
		if !ifInIdMap(key.(string)) {
			// not in
			return true
		}

		agent, ok := value.(*gocCoveredAgent)
		if !ok {
			return false
		}

		// check if extra matches
		if !isExtra(agent.Extra) {
			// not match
			return true
		}

		var req ProfileReq = "resetprofile"
		var res ProfileRes
		go func() {
			// lock-free
			rpc := agent.rpc
			if rpc == nil || agent.Status == DISCONNECT {
				return
			}
			err := rpc.Call("GocAgent.ResetProfile", req, &res)
			if err != nil {
				log.Errorf("fail to reset profile from: %v, reasson: %v. let's close the connection", agent.Id, err)
				// 关闭链接
				agent.closeRpcConnOnce()
			}
		}()

		return true
	})
}

// watchProfileUpdate watch the profile change
//
// any profile change will be updated on this websocket connection.
func (gs *gocServer) watchProfileUpdate(c *gin.Context) {
	// upgrade to websocket
	ws, err := gs.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Errorf("fail to establish websocket connection with watch client: %v", err)
		c.JSON(http.StatusInternalServerError, nil)
	}

	log.Infof("watch client connected")

	id := time.Now().String()
	gwc := &gocWatchClient{
		ws:     ws,
		exitCh: make(chan int),
	}
	gs.watchClients.Store(id, gwc)
	// send close msg and close ws connection
	defer func() {
		gs.watchClients.Delete(id)
		ws.Close()
		gwc.once.Do(func() { close(gwc.exitCh) })
		log.Infof("watch client disconnected")
	}()

	// set pong handler
	ws.SetReadDeadline(time.Now().Add(PongWait))
	ws.SetPongHandler(func(string) error {
		ws.SetReadDeadline(time.Now().Add(PongWait))
		return nil
	})

	// set ping goroutine to ping every PingWait time
	go func() {
		ticker := time.NewTicker(PingWait)
		defer ticker.Stop()

		for range ticker.C {
			if err := gs.wsping(ws, PongWait); err != nil {
				break
			}
		}

		gwc.once.Do(func() { close(gwc.exitCh) })
	}()

	<-gwc.exitCh
}

func filterProfileByPattern(skippattern []string, needpattern []string, profiles []*cover.Profile) []*cover.Profile {

	var out = make([]*cover.Profile, 0)
	var skipOut = make([]*cover.Profile, 0)

	if len(skippattern) != 0 {
		for _, profile := range profiles {
			skip := false
			for _, pattern := range skippattern {
				if strings.Contains(profile.FileName, pattern) {
					skip = true
					break
				}
			}

			if !skip {
				skipOut = append(skipOut, profile)
			}
		}
	} else {
		return profiles
	}
	log.Infof("skipOut len: %v", len(skipOut))
	if len(needpattern) == 0 {
		return skipOut
	}

	for _, profile := range skipOut {
		need := false
		for _, pattern := range needpattern {
			if strings.Contains(profile.FileName, pattern) {
				need = true
				break
			}
		}
		if need {
			out = append(out, profile)
		}
	}
	log.Infof("need out len: %v", len(out))

	return out
}

func idMaps(idQuery string) func(key string) bool {
	idMap := make(map[string]bool)
	if len(strings.TrimSpace(idQuery)) == 0 {
	} else {
		ids := strings.Split(idQuery, ",")
		for _, id := range ids {
			idMap[id] = true
		}
	}

	inIdMaps := func(key string) bool {
		// if no id in query, then all id agent will be return
		if len(idMap) == 0 {
			return true
		}
		// other
		_, ok := idMap[key]
		if !ok {
			return false
		} else {
			return true
		}
	}

	return inIdMaps
}

func filterExtra(extraPattern string) func(string) bool {

	re := regexp.MustCompile(extraPattern)

	return func(extra string) bool {
		return re.Match([]byte(extra))
	}
}