mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-09-20 13:30:46 +02:00
462 lines
12 KiB
Go
462 lines
12 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"sync"
|
|
|
|
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
|
|
"google.golang.org/protobuf/types/known/emptypb"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
// GrpcServer implements the plugin gRPC service
|
|
type GrpcServer struct {
|
|
plugin_pb.UnimplementedPluginServiceServer
|
|
plugin_pb.UnimplementedAdminQueryServiceServer
|
|
plugin_pb.UnimplementedAdminCommandServiceServer
|
|
|
|
registry *Registry
|
|
dispatcher *Dispatcher
|
|
configMgr *ConfigManager
|
|
|
|
// Active connections
|
|
connections map[string]*grpc.ClientStream
|
|
connMu sync.RWMutex
|
|
}
|
|
|
|
// NewGrpcServer creates a new gRPC server
|
|
func NewGrpcServer(registry *Registry, dispatcher *Dispatcher, configMgr *ConfigManager) *GrpcServer {
|
|
return &GrpcServer{
|
|
registry: registry,
|
|
dispatcher: dispatcher,
|
|
configMgr: configMgr,
|
|
connections: make(map[string]*grpc.ClientStream),
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// PluginService Implementation
|
|
// ============================================================================
|
|
|
|
// Connect handles bidirectional streaming for plugin registration and lifecycle
|
|
func (s *GrpcServer) Connect(stream grpc.BidiStreamingServer[plugin_pb.PluginMessage, plugin_pb.AdminMessage]) error {
|
|
var pluginID string
|
|
|
|
for {
|
|
msg, err := stream.Recv()
|
|
if err == io.EOF {
|
|
// Stream closed normally
|
|
if pluginID != "" {
|
|
s.registry.UnregisterPlugin(pluginID)
|
|
}
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Handle different message types
|
|
switch content := msg.Content.(type) {
|
|
case *plugin_pb.PluginMessage_Register:
|
|
register := content.Register
|
|
pluginID = register.PluginId
|
|
_, _ = s.registry.RegisterPlugin(pluginID, register)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.PluginMessage_Heartbeat:
|
|
heartbeat := content.Heartbeat
|
|
if err := s.registry.UpdateHeartbeat(heartbeat.PluginId, heartbeat); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.PluginMessage_StatusUpdate:
|
|
// Handle job status updates
|
|
statusUpdate := content.StatusUpdate
|
|
if err := s.handleStatusUpdate(statusUpdate); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Send acknowledgment
|
|
adminMsg := &plugin_pb.AdminMessage{
|
|
Content: &plugin_pb.AdminMessage_AdminCommand{
|
|
AdminCommand: &plugin_pb.AdminCommand{
|
|
CommandType: plugin_pb.AdminCommand_RELOAD_CONFIG,
|
|
},
|
|
},
|
|
}
|
|
|
|
if err := stream.Send(adminMsg); err != nil {
|
|
if pluginID != "" {
|
|
s.registry.UnregisterPlugin(pluginID)
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// ExecuteJob handles bidirectional streaming for job execution
|
|
func (s *GrpcServer) ExecuteJob(stream grpc.BidiStreamingServer[plugin_pb.JobExecutionMessage, plugin_pb.JobProgressMessage]) error {
|
|
var jobID string
|
|
|
|
for {
|
|
execMsg, err := stream.Recv()
|
|
if err == io.EOF {
|
|
// Stream closed normally
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
jobID = execMsg.JobId
|
|
|
|
// Handle different execution message types
|
|
switch content := execMsg.Content.(type) {
|
|
case *plugin_pb.JobExecutionMessage_JobStarted:
|
|
started := content.JobStarted
|
|
if err := s.handleJobStarted(jobID, started); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.JobExecutionMessage_Progress:
|
|
progress := content.Progress
|
|
if err := s.dispatcher.UpdateJobProgress(jobID, progress.ProgressPercent); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.JobExecutionMessage_Checkpoint:
|
|
checkpoint := content.Checkpoint
|
|
// Store checkpoint for resumption
|
|
if err := s.storeCheckpoint(jobID, checkpoint); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.JobExecutionMessage_JobCompleted:
|
|
completed := content.JobCompleted
|
|
if err := s.dispatcher.CompleteJob(jobID, completed); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.JobExecutionMessage_JobFailed:
|
|
failed := content.JobFailed
|
|
if err := s.dispatcher.FailJob(jobID, failed); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *plugin_pb.JobExecutionMessage_LogEntry:
|
|
logEntry := content.LogEntry
|
|
if err := s.storeLogEntry(jobID, logEntry); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Send progress update
|
|
progressMsg := &plugin_pb.JobProgressMessage{
|
|
JobId: jobID,
|
|
Content: &plugin_pb.JobProgressMessage_Update{
|
|
Update: &plugin_pb.JobProgressUpdate{
|
|
JobId: jobID,
|
|
ProgressPercent: 0,
|
|
CurrentStep: "processing",
|
|
StatusMessage: "Processing job",
|
|
Timestamp: nil,
|
|
},
|
|
},
|
|
}
|
|
|
|
if err := stream.Send(progressMsg); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// AdminQueryService Implementation
|
|
// ============================================================================
|
|
|
|
// GetPluginStats returns current plugin system statistics
|
|
func (s *GrpcServer) GetPluginStats(ctx context.Context, req *emptypb.Empty) (*plugin_pb.PluginStats, error) {
|
|
plugins := s.registry.ListPlugins()
|
|
jobs := s.dispatcher.ListAllJobs()
|
|
|
|
var totalJobs, pendingJobs, runningJobs, completedJobs, failedJobs int32
|
|
|
|
for _, job := range jobs {
|
|
totalJobs++
|
|
switch job.GetState() {
|
|
case JobStatePending:
|
|
pendingJobs++
|
|
case JobStateRunning:
|
|
runningJobs++
|
|
case JobStateCompleted:
|
|
completedJobs++
|
|
case JobStateFailed:
|
|
failedJobs++
|
|
}
|
|
}
|
|
|
|
stats := &plugin_pb.PluginStats{
|
|
TotalPlugins: int32(len(plugins)),
|
|
ActivePlugins: int32(len(plugins)),
|
|
TotalJobs: totalJobs,
|
|
PendingJobs: pendingJobs,
|
|
RunningJobs: runningJobs,
|
|
CompletedJobs: completedJobs,
|
|
FailedJobs: failedJobs,
|
|
}
|
|
|
|
return stats, nil
|
|
}
|
|
|
|
// ListPlugins returns all connected plugins
|
|
func (s *GrpcServer) ListPlugins(ctx context.Context, req *emptypb.Empty) (*plugin_pb.PluginList, error) {
|
|
plugins := s.registry.ListPlugins()
|
|
pluginInfos := make([]*plugin_pb.PluginInfo, 0, len(plugins))
|
|
|
|
for _, p := range plugins {
|
|
info := &plugin_pb.PluginInfo{
|
|
PluginId: p.ID,
|
|
Name: p.Name,
|
|
Version: p.Version,
|
|
ProtocolVersion: p.ProtocolVersion,
|
|
Healthy: p.IsHealthy(),
|
|
Status: p.GetState().String(),
|
|
}
|
|
|
|
// Add capabilities
|
|
for _, cap := range p.Capabilities {
|
|
info.Capabilities = append(info.Capabilities, cap)
|
|
}
|
|
|
|
pluginInfos = append(pluginInfos, info)
|
|
}
|
|
|
|
return &plugin_pb.PluginList{Plugins: pluginInfos}, nil
|
|
}
|
|
|
|
// ListJobs returns jobs filtered by type and status
|
|
func (s *GrpcServer) ListJobs(ctx context.Context, req *plugin_pb.ListJobsRequest) (*plugin_pb.JobList, error) {
|
|
var state JobState
|
|
switch req.Status {
|
|
case plugin_pb.ListJobsRequest_PENDING:
|
|
state = JobStatePending
|
|
case plugin_pb.ListJobsRequest_RUNNING:
|
|
state = JobStateRunning
|
|
case plugin_pb.ListJobsRequest_COMPLETED:
|
|
state = JobStateCompleted
|
|
case plugin_pb.ListJobsRequest_FAILED:
|
|
state = JobStateFailed
|
|
default:
|
|
state = JobStatePending
|
|
}
|
|
|
|
jobs := s.dispatcher.ListJobs(req.JobType, state)
|
|
|
|
jobPbs := make([]*plugin_pb.Job, 0, len(jobs))
|
|
for _, job := range jobs {
|
|
jobPb := &plugin_pb.Job{
|
|
JobId: job.ID,
|
|
JobType: job.Type,
|
|
Description: job.Description,
|
|
Priority: job.Priority,
|
|
CreatedAt: nil,
|
|
UpdatedAt: nil,
|
|
ExecutorPluginId: job.ExecutorID,
|
|
ProgressPercent: job.GetProgress(),
|
|
Config: job.Config,
|
|
}
|
|
jobPbs = append(jobPbs, jobPb)
|
|
}
|
|
|
|
return &plugin_pb.JobList{
|
|
Jobs: jobPbs,
|
|
TotalCount: int32(len(jobPbs)),
|
|
ReturnedCount: int32(len(jobPbs)),
|
|
}, nil
|
|
}
|
|
|
|
// GetJob returns detailed job information
|
|
func (s *GrpcServer) GetJob(ctx context.Context, req *plugin_pb.GetJobRequest) (*plugin_pb.JobDetail, error) {
|
|
job, err := s.dispatcher.GetJob(req.JobId)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("job not found: %w", err)
|
|
}
|
|
|
|
jobPb := &plugin_pb.Job{
|
|
JobId: job.ID,
|
|
JobType: job.Type,
|
|
Description: job.Description,
|
|
Priority: job.Priority,
|
|
ExecutorPluginId: job.ExecutorID,
|
|
ProgressPercent: job.GetProgress(),
|
|
Config: job.Config,
|
|
}
|
|
|
|
detail := &plugin_pb.JobDetail{
|
|
Job: jobPb,
|
|
LogEntries: job.LogEntries,
|
|
CompletionInfo: job.CompletedInfo,
|
|
FailureInfo: job.FailedInfo,
|
|
CheckpointIds: job.CheckpointIDs,
|
|
}
|
|
|
|
return detail, nil
|
|
}
|
|
|
|
// GetJobExecutionLog returns job execution logs
|
|
func (s *GrpcServer) GetJobExecutionLog(req *plugin_pb.GetJobExecutionLogRequest, stream grpc.ServerStreamingServer[plugin_pb.ExecutionLogEntry]) error {
|
|
job, err := s.dispatcher.GetJob(req.JobId)
|
|
if err != nil {
|
|
return fmt.Errorf("job not found: %w", err)
|
|
}
|
|
|
|
// Send log entries
|
|
for _, entry := range job.LogEntries {
|
|
if err := stream.Send(entry); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ============================================================================
|
|
// AdminCommandService Implementation
|
|
// ============================================================================
|
|
|
|
// UpdateJobTypeConfig updates configuration for a job type
|
|
func (s *GrpcServer) UpdateJobTypeConfig(ctx context.Context, req *plugin_pb.UpdateJobTypeConfigRequest) (*plugin_pb.JobTypeConfig, error) {
|
|
config := req.Config
|
|
|
|
_, err := s.configMgr.UpdateConfig(
|
|
config.JobType,
|
|
config.Enabled,
|
|
config.AdminConfig,
|
|
config.WorkerConfig,
|
|
"admin",
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to update config: %w", err)
|
|
}
|
|
|
|
return config, nil
|
|
}
|
|
|
|
// CreateJob creates a new job
|
|
func (s *GrpcServer) CreateJob(ctx context.Context, req *plugin_pb.CreateJobRequest) (*plugin_pb.Job, error) {
|
|
jobReq := &plugin_pb.JobRequest{
|
|
JobType: req.JobType,
|
|
Description: req.Description,
|
|
Priority: req.Priority,
|
|
Config: req.Config,
|
|
Metadata: req.Metadata,
|
|
RequestSource: "admin",
|
|
}
|
|
|
|
job, err := s.dispatcher.QueueJob(req.JobType, jobReq, "")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create job: %w", err)
|
|
}
|
|
|
|
return &plugin_pb.Job{
|
|
JobId: job.ID,
|
|
JobType: job.Type,
|
|
Description: job.Description,
|
|
Priority: job.Priority,
|
|
ProgressPercent: job.GetProgress(),
|
|
Config: job.Config,
|
|
}, nil
|
|
}
|
|
|
|
// CancelJob cancels a running job
|
|
func (s *GrpcServer) CancelJob(ctx context.Context, req *plugin_pb.CancelJobRequest) (*emptypb.Empty, error) {
|
|
if err := s.dispatcher.CancelJob(req.JobId); err != nil {
|
|
return nil, fmt.Errorf("failed to cancel job: %w", err)
|
|
}
|
|
|
|
return &emptypb.Empty{}, nil
|
|
}
|
|
|
|
// RetryJob retries a failed job
|
|
func (s *GrpcServer) RetryJob(ctx context.Context, req *plugin_pb.RetryJobRequest) (*plugin_pb.Job, error) {
|
|
if err := s.dispatcher.RetryJob(req.JobId); err != nil {
|
|
return nil, fmt.Errorf("failed to retry job: %w", err)
|
|
}
|
|
|
|
job, err := s.dispatcher.GetJob(req.JobId)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("job not found: %w", err)
|
|
}
|
|
|
|
return &plugin_pb.Job{
|
|
JobId: job.ID,
|
|
JobType: job.Type,
|
|
Description: job.Description,
|
|
Priority: job.Priority,
|
|
ProgressPercent: job.GetProgress(),
|
|
Config: job.Config,
|
|
}, nil
|
|
}
|
|
|
|
// ============================================================================
|
|
// Helper Methods
|
|
// ============================================================================
|
|
|
|
// handleStatusUpdate handles execution status updates
|
|
func (s *GrpcServer) handleStatusUpdate(update *plugin_pb.ExecutionStatusUpdate) error {
|
|
if err := s.dispatcher.UpdateJobProgress(update.JobId, update.ProgressPercent); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// handleJobStarted handles job start messages
|
|
func (s *GrpcServer) handleJobStarted(jobID string, started *plugin_pb.JobStarted) error {
|
|
job, err := s.dispatcher.GetJob(jobID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
job.SetState(JobStateRunning)
|
|
return nil
|
|
}
|
|
|
|
// storeCheckpoint stores a job checkpoint
|
|
func (s *GrpcServer) storeCheckpoint(jobID string, checkpoint *plugin_pb.JobCheckpoint) error {
|
|
job, err := s.dispatcher.GetJob(jobID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
job.mu.Lock()
|
|
defer job.mu.Unlock()
|
|
|
|
job.CheckpointIDs = append(job.CheckpointIDs, checkpoint.CheckpointId)
|
|
return nil
|
|
}
|
|
|
|
// storeLogEntry stores a log entry for a job
|
|
func (s *GrpcServer) storeLogEntry(jobID string, logEntry *plugin_pb.ExecutionLog) error {
|
|
job, err := s.dispatcher.GetJob(jobID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
job.mu.Lock()
|
|
defer job.mu.Unlock()
|
|
|
|
job.LogEntries = append(job.LogEntries, &plugin_pb.ExecutionLogEntry{
|
|
Timestamp: logEntry.Timestamp,
|
|
Level: logEntry.Level,
|
|
Message: logEntry.Message,
|
|
Context: logEntry.Context,
|
|
})
|
|
|
|
return nil
|
|
}
|