threshold/internal/servers/worker/status.go

package worker

import (
	"context"
	"math/rand"
	"time"

	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
	"github.com/hashicorp/boundary/internal/servers"
	"github.com/hashicorp/boundary/internal/types/resource"
	"google.golang.org/grpc/resolver"
)

// In the future we could make this configurable
const (
	statusInterval = 2 * time.Second
)

type LastStatusInformation struct {
	*pbs.StatusResponse
	StatusTime time.Time
}

func (w *Worker) startStatusTicking(cancelCtx context.Context) {
	go func() {
		r := rand.New(rand.NewSource(time.Now().UnixNano()))
		// This function exists to desynchronize calls to controllers from
		// workers so we aren't always getting status updates at the exact same
		// intervals, to ease the load on the DB.
		getRandomInterval := func() time.Duration {
			// 0 to 0.5 adjustment to the base
			f := r.Float64() / 2
			// Half a chance to be faster, not slower
			if r.Float32() > 0.5 {
				f = -1 * f
			}
			return statusInterval + time.Duration(f*float64(time.Second))
		}

		timer := time.NewTimer(0)
		for {
			select {
			case <-cancelCtx.Done():
				w.logger.Info("status ticking shutting down")
				return

			case <-timer.C:
				w.sendWorkerStatus(cancelCtx)
				timer.Reset(getRandomInterval())
			}
		}
	}()
}

// LastStatusSuccess is used in tests (it's exported for tests in other
// packages) to verify the last time we sent status
func (w *Worker) LastStatusSuccess() *LastStatusInformation {
	return w.lastStatusSuccess.Load().(*LastStatusInformation)
}

func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
	// First send info as-is. We'll perform cleanup duties after we
	// get cancel/job change info back.
	var activeJobs []*pbs.JobStatus

	// Range over known sessions and collect info
	w.sessionInfoMap.Range(func(key, value interface{}) bool {
		var jobInfo pbs.SessionJobInfo
		sessionId := key.(string)
		si := value.(*sessionInfo)
		si.RLock()
		status := si.status
		connections := make([]*pbs.Connection, 0, len(si.connInfoMap))
		for k, v := range si.connInfoMap {
			connections = append(connections, &pbs.Connection{
				ConnectionId: k,
				Status:       v.status,
			})
		}
		si.RUnlock()
		jobInfo.SessionId = sessionId
		activeJobs = append(activeJobs, &pbs.JobStatus{
			Job: &pbs.Job{
				Type: pbs.JOBTYPE_JOBTYPE_SESSION,
				JobInfo: &pbs.Job_SessionInfo{
					SessionInfo: &pbs.SessionJobInfo{
						SessionId:   sessionId,
						Status:      status,
						Connections: connections,
					},
				},
			},
		})
		return true
	})

	// Send status information
	client := w.controllerStatusConn.Load().(pbs.ServerCoordinationServiceClient)
	var tags map[string]*servers.TagValues
	// If we're not going to request a tag update, no reason to have these
	// marshaled on every status call.
	if w.updateTags.Load() {
		tags = w.tags.Load().(map[string]*servers.TagValues)
	}
	result, err := client.Status(cancelCtx, &pbs.StatusRequest{
		Jobs: activeJobs,
		Worker: &servers.Server{
			PrivateId:   w.conf.RawConfig.Worker.Name,
			Type:        resource.Worker.String(),
			Description: w.conf.RawConfig.Worker.Description,
			Address:     w.conf.RawConfig.Worker.PublicAddr,
			Tags:        tags,
		},
		UpdateTags: w.updateTags.Load(),
	})
	if err != nil {
		w.logger.Error("error making status request to controller", "error", err)
	} else {
		w.logger.Trace("successfully sent status to controller")
		w.updateTags.Store(false)
		addrs := make([]resolver.Address, 0, len(result.Controllers))
		strAddrs := make([]string, 0, len(result.Controllers))
		for _, v := range result.Controllers {
			addrs = append(addrs, resolver.Address{Addr: v.Address})
			strAddrs = append(strAddrs, v.Address)
		}
		w.logger.Trace("found controllers", "addresses", strAddrs)
		switch len(strAddrs) {
		case 0:
			w.logger.Warn("got no controller addresses from controller; possibly prior to first status save, not persisting")
		default:
			w.Resolver().UpdateState(resolver.State{Addresses: addrs})
		}
		w.lastStatusSuccess.Store(&LastStatusInformation{StatusResponse: result, StatusTime: time.Now()})

		for _, request := range result.GetJobsRequests() {
			switch request.GetRequestType() {
			case pbs.CHANGETYPE_CHANGETYPE_UPDATE_STATE:
				switch request.GetJob().GetType() {
				case pbs.JOBTYPE_JOBTYPE_SESSION:
					sessInfo := request.GetJob().GetSessionInfo()
					sessionId := sessInfo.GetSessionId()
					siRaw, ok := w.sessionInfoMap.Load(sessionId)
					if !ok {
						w.logger.Warn("asked to cancel session but could not find a local information for it", "session_id", sessionId)
						continue
					}
					si := siRaw.(*sessionInfo)
					si.Lock()
					si.status = sessInfo.GetStatus()
					si.Unlock()
				}
			}
		}
	}

	// Cleanup: Run through current jobs. Cancel connections for any
	// canceling session or any session that is expired. Clear out
	// sessions that are canceled or expired with all connections
	// marked as closed. Close any that aren't marked as such.
	closeInfo := make(map[string]string)
	cleanSessionIds := make([]string, 0)
	w.sessionInfoMap.Range(func(key, value interface{}) bool {
		si := value.(*sessionInfo)
		si.Lock()
		switch {
		case si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
			si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
			time.Until(si.lookupSessionResponse.Expiration.AsTime()) < 0:
			var toClose int
			for k, v := range si.connInfoMap {
				if v.closeTime.IsZero() {
					toClose++
					v.connCancel()
					w.logger.Info("terminated connection due to cancelation or expiration", "session_id", si.id, "connection_id", k)
					closeInfo[k] = si.id
				}
			}
			// closeTime is marked by closeConnections iff the
			// status is returned for that connection as closed. If
			// the session is no longer valid and all connections
			// are marked closed, clean up the session.
			if toClose == 0 {
				cleanSessionIds = append(cleanSessionIds, si.id)
			}
		}
		si.Unlock()
		return true
	})

	// Note that we won't clean these from the info map until the
	// next time we run this function
	if len(closeInfo) > 0 {
		if err := w.closeConnections(cancelCtx, closeInfo); err != nil {
			w.logger.Error("error marking connections closed", "error", err)
		}
	}

	// Forget sessions where the session is expired/canceled and all
	// connections are canceled and marked closed
	for _, v := range cleanSessionIds {
		w.sessionInfoMap.Delete(v)
	}
}