HellreigN/backend/internal/grpcsrv/commander/commander.go

package commander

import (
	"context"
	"fmt"
	"io"
	"sync"

	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/models"
	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/proto/proto"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type Commander struct {
	proto.UnimplementedCommanderServer
	agents map[string]Agent
	mu     sync.RWMutex
	jobber Jobber
}

type Jobber interface {
	InitJob(ctx context.Context, agentID string, job models.JobForInsert) (int64, error)
	UpdateJobInDB(ctx context.Context, jid int64, msg models.JobForUpdate) (models.Job, error)
}

func New(jobber Jobber) *Commander {
	return &Commander{
		agents: make(map[string]Agent),
		jobber: jobber,
	}
}

type Agent struct {
	bidi   grpc.BidiStreamingServer[proto.FinishedCommand, proto.Command]
	in     chan *proto.Command
	jobs   map[int64]Job
	jobber Jobber
	ctx    context.Context
	aid    string

	Token    string // agent id
	Label    string
	Services []string
}
type JobOut struct {
	fc  models.Job
	err error
}

type Job struct {
	out chan JobOut
}

func (self *Commander) GetAgent(aid string) (agent Agent, ok bool) {
	self.mu.RLock()
	defer self.mu.RUnlock()
	agent, ok = self.agents[aid]
	return
}

func (self *Commander) Agents() []Agent {
	self.mu.RLock()
	defer self.mu.RUnlock()
	result := make([]Agent, 0, len(self.agents))
	for _, a := range self.agents {
		result = append(result, a)
	}
	return result
}

func (self *Commander) removeAgent(aid string) {
	self.mu.Lock()
	defer self.mu.Unlock()
	delete(self.agents, aid)
}

func (self *Agent) AddJob(job models.JobForInsert) (int64, error) {
	jid, err := self.jobber.InitJob(self.ctx, self.aid, job)
	if err != nil {
		return 0, err
	}
	self.jobs[jid] = newJob()
	self.in <- &proto.Command{
		Id:      jid,
		Command: job.Command,
		Stdin:   job.Stdin,
	}
	return jid, err
}

func (self *Agent) WaitJob(jid int64) (*models.Job, error) {
	result := <-self.jobs[jid].out
	return &result.fc, result.err
}

func (self *Commander) Stream(bidi grpc.BidiStreamingServer[proto.FinishedCommand, proto.Command]) error {
	md, ok := metadata.FromIncomingContext(bidi.Context())
	if !ok {
		return fmt.Errorf("no metadata in context")
	}
	aidVals := md["agentid"]
	if len(aidVals) == 0 {
		return fmt.Errorf("agentid metadata missing")
	}
	aid := aidVals[0]

	var label string
	labelVals := md["label"]
	if len(labelVals) > 0 {
		label = labelVals[0]
	}

	agent := newAgent(bidi, self.jobber, aid, label)
	self.mu.Lock()
	self.agents[aid] = agent
	self.mu.Unlock()

	defer self.removeAgent(aid)
	return agent.run()
}

func (self *Agent) run() error {
	wg := new(errgroup.Group)
	wg.Go(self.recv)
	wg.Go(self.send)
	return wg.Wait()
}

func (self *Agent) recv() error {
	for {
		job, err := func() (job models.Job, err error) {
			msg, err := self.bidi.Recv()
			if err != nil {
				return
			}
			return self.jobber.UpdateJobInDB(self.ctx, msg.Id, models.JobForUpdate{
				Stdout: msg.Stdout,
				Stderr: msg.Stderr,
				Status: msg.Status,
			})
		}()
		if err == io.EOF {
			return nil
		}
		// TODO: that would blow up at some point
		out := self.jobs[job.ID].out
		out <- JobOut{
			fc:  job,
			err: err,
		}
		close(out)
	}
}

func (self *Agent) send() error {
	for job := range self.in {
		if err := self.bidi.Send(job); err != nil {
			return err
		}
	}
	return io.EOF
	// self.jobs[]
}

func newAgent(bidi grpc.BidiStreamingServer[proto.FinishedCommand, proto.Command], jobber Jobber, aid string, label string) Agent {
	return Agent{
		bidi:   bidi,
		in:     make(chan *proto.Command),
		jobs:   make(map[int64]Job),
		jobber: jobber,
		ctx:    bidi.Context(),
		aid:    aid,
		Label:  label,
		Token:  aid,
	}
}

func newJob() Job {
	return Job{make(chan JobOut, 1)}
}