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 && len(needpattern) == 0 {
        return profiles
    }
    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 {
        skipOut = 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))
    }
}