okooooooo/pika
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support disabling agents

Browse Source
This commit is contained in:
dushixiang committed 2026-08-26 20:24:17 +08:00
1 parent a4ec00b47e
commit 14a48f81a4
13 files changed
+309 -10
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/app.go
+2
View File
@@ -208,6 +208,8 @@ func setupApi(app *orz.App, components *AppComponents) error {
adminApi.GET("/agents/:id", components.AgentHandler.GetForAdmin)
adminApi.GET("/agents/:id/metrics/latest", components.AgentHandler.GetAdminLatestMetrics)
adminApi.PUT("/agents/:id", components.AgentHandler.UpdateInfo)
adminApi.POST("/agents/:id/enable", components.AgentHandler.Enable)
adminApi.POST("/agents/:id/disable", components.AgentHandler.Disable)
adminApi.POST("/agents/batch/tags", components.AgentHandler.BatchUpdateTags)
adminApi.POST("/agents/batch/visibility", components.AgentHandler.BatchUpdateVisibility)
adminApi.DELETE("/agents/:id", components.AgentHandler.Delete)
internal/handler/agent_enabled_test.go
+81
View File
@@ -0,0 +1,81 @@
package handler
import (
"context"
"encoding/json"
"testing"
"github.com/glebarez/sqlite"
"github.com/pika-monitor/pika/internal/models"
"github.com/pika-monitor/pika/internal/protocol"
"github.com/pika-monitor/pika/internal/repo"
"github.com/pika-monitor/pika/internal/service"
"go.uber.org/zap"
"gorm.io/gorm"
)
func TestHandleWebSocketMessageIgnoresDisabledAgent(t *testing.T) {
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("open database: %v", err)
}
if err := database.AutoMigrate(&models.Agent{}); err != nil {
t.Fatalf("migrate agent: %v", err)
}
agent := models.Agent{ID: "agent-disabled", Enabled: true}
if err := database.Create(&agent).Error; err != nil {
t.Fatalf("create agent: %v", err)
}
if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", false).Error; err != nil {
t.Fatalf("disable agent: %v", err)
}
h := &AgentHandler{
logger: zap.NewNop(),
agentService: &service.AgentService{
AgentRepo: repo.NewAgentRepo(database),
},
}
invalidMetrics := json.RawMessage(`{`)
if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err != nil {
t.Fatalf("disabled agent data should be ignored, got error: %v", err)
}
if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", true).Error; err != nil {
t.Fatalf("enable agent: %v", err)
}
if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err == nil {
t.Fatal("enabled agent data should reach the metrics parser")
}
}
func TestDisabledAgentStatusCannotBeRefreshed(t *testing.T) {
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("open database: %v", err)
}
if err := database.AutoMigrate(&models.Agent{}); err != nil {
t.Fatalf("migrate agent: %v", err)
}
agent := models.Agent{ID: "agent-status", Enabled: true, Status: 1}
if err := database.Create(&agent).Error; err != nil {
t.Fatalf("create agent: %v", err)
}
agentService := &service.AgentService{AgentRepo: repo.NewAgentRepo(database)}
if err := agentService.UpdateAgentEnabled(context.Background(), agent.ID, false); err != nil {
t.Fatalf("disable agent: %v", err)
}
if err := agentService.UpdateAgentStatus(context.Background(), agent.ID, 1); err != nil {
t.Fatalf("refresh disabled agent status: %v", err)
}
var saved models.Agent
if err := database.First(&saved, "id = ?", agent.ID).Error; err != nil {
t.Fatalf("load agent: %v", err)
}
if saved.Enabled || saved.Status != 0 {
t.Fatalf("disabled agent status changed unexpectedly: enabled=%v status=%d", saved.Enabled, saved.Status)
}
}
internal/handler/agent_handler.go
+26
View File
@@ -1,13 +1,39 @@
package handler
import (
"sync"
"github.com/go-orz/orz"
"github.com/gorilla/websocket"
"github.com/labstack/echo/v5"
"github.com/pika-monitor/pika/internal/service"
ws "github.com/pika-monitor/pika/internal/websocket"
"go.uber.org/zap"
)
// Enable 启用探针数据处理。
func (h *AgentHandler) Enable(c *echo.Context) error {
return h.updateEnabled(c, true)
}
// Disable 禁用探针数据处理和告警。
func (h *AgentHandler) Disable(c *echo.Context) error {
return h.updateEnabled(c, false)
}
func (h *AgentHandler) updateEnabled(c *echo.Context, enabled bool) error {
h.enabledMu.Lock()
defer h.enabledMu.Unlock()
agentID := c.Param("id")
if err := h.agentService.UpdateAgentEnabled(c.Request().Context(), agentID, enabled); err != nil {
return err
}
return orz.Ok(c, orz.Map{})
}
type AgentHandler struct {
enabledMu sync.RWMutex
logger *zap.Logger
agentService *service.AgentService
trafficService *service.TrafficService
internal/handler/agent_query_handler.go
+15
View File
@@ -1,6 +1,8 @@
package handler
import (
"encoding/json"
"github.com/go-orz/orz"
"github.com/labstack/echo/v5"
"github.com/pika-monitor/pika/internal/models"
@@ -32,6 +34,19 @@ func (h *AgentHandler) Get(c *echo.Context) error {
agent.Hostname = ""
}
if !isAuthenticated {
payload, err := json.Marshal(agent)
if err != nil {
return err
}
var publicAgent map[string]any
if err := json.Unmarshal(payload, &publicAgent); err != nil {
return err
}
delete(publicAgent, "enabled")
return orz.Ok(c, publicAgent)
}
return orz.Ok(c, agent)
}
internal/handler/agent_ws_handler.go
+15
View File
@@ -28,7 +28,9 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
}
// 注册探针 - 使用独立的context,不依赖HTTP请求的context
h.enabledMu.RLock()
agent, err := h.agentService.RegisterAgent(context.Background(), c.RealIP(), &registerReq.AgentInfo, registerReq.ApiKey)
h.enabledMu.RUnlock()
if err != nil {
// 发送注册失败响应
h.sendRegisterError(conn, err.Error())
@@ -47,6 +49,7 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
return err
}
if agent.Enabled {
// 下发防篡改配置
if err := h.sendTamperConfig(conn, agent.ID); err != nil {
h.logger.Error("failed to send tamper config", zap.Error(err))
@@ -62,6 +65,7 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
h.logger.Error("failed to send public ip config", zap.Error(err))
// 配置下发失败不中断连接，只记录日志
}
}
// 创建客户端并注册到管理器
client := h.newClient(agent.ID, conn)
@@ -76,6 +80,17 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
// handleWebSocketMessage 处理WebSocket消息
func (h *AgentHandler) handleWebSocketMessage(ctx context.Context, agentID string, messageType string, data json.RawMessage) error {
h.enabledMu.RLock()
defer h.enabledMu.RUnlock()
enabled, err := h.agentService.IsAgentEnabled(ctx, agentID)
if err != nil {
return err
}
if !enabled {
return nil
}
switch protocol.MessageType(messageType) {
case protocol.MessageTypeMetrics:
return h.handleMetricsMessage(ctx, agentID, data)
internal/models/agent.go
+1
View File
@@ -20,6 +20,7 @@ type Agent struct {
Version string `json:"version"` // 探针版本
Tags datatypes.JSONSlice[string] `json:"tags"` // 标签
ExpireTime int64 `json:"expireTime"` // 到期时间（时间戳毫秒）
Enabled bool `gorm:"default:true;not null;index" json:"enabled"` // 是否接收并处理探针数据
Status int `json:"status"` // 状态: 0-离线, 1-在线
Visibility string `gorm:"default:public" json:"visibility"` // 可见性: public-匿名可见, private-登录可见
Weight int `gorm:"default:0;index" json:"weight"` // 展示顺序（由拖动排序维护，数字越大越靠前）
internal/repo/agent_repo.go
+15 -3
View File
@@ -32,15 +32,27 @@ func (r *AgentRepo) UpdateStatus(ctx context.Context, agentID string, status int
return r.db.WithContext(ctx).
Model(&models.Agent{}).
Where("id = ?", agentID).
Where("id = ? AND enabled = ?", agentID, true).
Updates(m).Error
}
// IsEnabled 查询探针是否允许接收并处理数据。
func (r *AgentRepo) IsEnabled(ctx context.Context, agentID string) (bool, error) {
var agent models.Agent
if err := r.db.WithContext(ctx).
Select("id", "enabled").
Where("id = ?", agentID).
First(&agent).Error; err != nil {
return false, err
}
return agent.Enabled, nil
}
// FindOnlineAgents 查找所有在线探针
func (r *AgentRepo) FindOnlineAgents(ctx context.Context) ([]models.Agent, error) {
var agents []models.Agent
err := r.db.WithContext(ctx).
Where("status = ?", 1).
Where("status = ? AND enabled = ?", 1, true).
Find(&agents).Error
return agents, err
}
@@ -136,7 +148,7 @@ func (r *AgentRepo) GetStatistics(ctx context.Context) (total int64, online int6
// 获取在线数量
err = r.db.WithContext(ctx).
Model(&models.Agent{}).
Where("status = ?", 1).
Where("status = ? AND enabled = ?", 1, true).
Count(&online).Error
return total, online, err
internal/service/agent_service.go
+30
View File
@@ -66,6 +66,14 @@ func (s *AgentService) RegisterAgent(ctx context.Context, ip string, info *proto
// 这样即使主机名变化，也能正确识别
existingAgent, err := s.AgentRepo.FindById(ctx, info.ID)
if err == nil {
// 已禁用的探针仅维持连接，不用注册报文刷新任何主机数据或在线状态。
if !existingAgent.Enabled {
s.logger.Info("disabled agent reconnected; registration data ignored",
zap.String("agentID", existingAgent.ID),
zap.String("hostname", existingAgent.Hostname))
return &existingAgent, nil
}
// 更新现有探针信息（允许主机名、名称等变化）
now := time.Now().UnixMilli()
existingAgent.Hostname = info.Hostname
@@ -117,6 +125,7 @@ func (s *AgentService) RegisterAgent(ctx context.Context, ip string, info *proto
OS: info.OS,
Arch: info.Arch,
Version: info.Version,
Enabled: true,
Status: 1,
LastSeenAt: now,
CreatedAt: now,
@@ -141,6 +150,27 @@ func (s *AgentService) UpdateAgentStatus(ctx context.Context, agentID string, st
return s.AgentRepo.UpdateStatus(ctx, agentID, status, time.Now().UnixMilli())
}
// IsAgentEnabled 查询探针是否允许接收并处理数据。
func (s *AgentService) IsAgentEnabled(ctx context.Context, agentID string) (bool, error) {
return s.AgentRepo.IsEnabled(ctx, agentID)
}
// UpdateAgentEnabled 更新探针启用状态。禁用时立即标记离线，后续状态更新由仓储层原子拦截。
func (s *AgentService) UpdateAgentEnabled(ctx context.Context, agentID string, enabled bool) error {
if _, err := s.AgentRepo.FindById(ctx, agentID); err != nil {
return err
}
updates := map[string]interface{}{
"enabled": enabled,
"updated_at": time.Now().UnixMilli(),
}
if !enabled {
updates["status"] = 0
}
return s.AgentRepo.UpdateColumnsById(ctx, agentID, updates)
}
// UpdatePublicIP 更新探针的公网 IP 信息
func (s *AgentService) UpdatePublicIP(ctx context.Context, agentID string, ipv4 string, ipv6 string) error {
updates := map[string]interface{}{
internal/service/alert_rule_disabled_agent_test.go
+55
View File
@@ -0,0 +1,55 @@
package service
import (
"context"
"testing"
"github.com/glebarez/sqlite"
"github.com/pika-monitor/pika/internal/models"
"go.uber.org/zap"
"gorm.io/gorm"
)
func TestResolveForAgentsExcludesDisabledAgents(t *testing.T) {
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("open database: %v", err)
}
if err := database.AutoMigrate(&models.Agent{}, &models.AlertRule{}); err != nil {
t.Fatalf("migrate models: %v", err)
}
agents := []models.Agent{
{ID: "enabled-agent", Enabled: true},
{ID: "disabled-agent", Enabled: true},
}
if err := database.Create(&agents).Error; err != nil {
t.Fatalf("create agents: %v", err)
}
if err := database.Model(&models.Agent{}).Where("id = ?", "disabled-agent").Update("enabled", false).Error; err != nil {
t.Fatalf("disable agent: %v", err)
}
rule := models.AlertRule{
ID: "all-agents-rule",
Name: "all agents",
Priority: 1,
Enabled: true,
TargetType: models.AlertRuleTargetAll,
}
if err := database.Create(&rule).Error; err != nil {
t.Fatalf("create alert rule: %v", err)
}
alertRuleService := NewAlertRuleService(zap.NewNop(), database, nil)
configs, err := alertRuleService.ResolveForAgents(context.Background(), []string{"enabled-agent", "disabled-agent"})
if err != nil {
t.Fatalf("resolve alert rules: %v", err)
}
if configs["enabled-agent"] == nil {
t.Fatal("enabled agent should match the alert rule")
}
if configs["disabled-agent"] != nil {
t.Fatal("disabled agent must not match any alert rule")
}
}
internal/service/alert_rule_service.go
+10 -1
View File
@@ -229,17 +229,26 @@ func (s *AlertRuleService) ResolveForAgents(ctx context.Context, agentIDs []stri
return result, nil
}
// 查询主机信息用于标签匹配
// 查询主机信息用于启用状态检查和标签匹配。禁用主机不命中任何告警规则，
// 从而统一阻止指标、离线、到期和监控等所有类型的告警。
agentTagMap := make(map[string][]string, len(agentIDs))
enabledAgents := make(map[string]struct{}, len(agentIDs))
agents, err := s.agentRepo.FindByIdIn(ctx, agentIDs)
if err != nil {
return nil, err
}
for _, agent := range agents {
if !agent.Enabled {
continue
}
enabledAgents[agent.ID] = struct{}{}
agentTagMap[agent.ID] = agent.Tags
}
for _, agentID := range agentIDs {
if _, enabled := enabledAgents[agentID]; !enabled {
continue
}
for i := range rules {
if !ruleMatchesAgent(&rules[i], agentID, agentTagMap[agentID]) {
continue
web/src/api/agent.ts
+5
View File
@@ -483,6 +483,11 @@ export const deleteAgent = (agentId: string) => {
return del(`/admin/agents/${agentId}`);
};
// 启用或禁用探针数据处理
export const updateAgentEnabled = (agentId: string, enabled: boolean) => {
return post(`/admin/agents/${agentId}/${enabled ? 'enable' : 'disable'}`);
};
// 获取所有探针的标签
export interface GetTagsResponse {
tags: string[];
web/src/pages/Agents/AgentList.tsx
+53 -6
View File
@@ -4,7 +4,7 @@ import {Link, useNavigate} from 'react-router-dom';
import type {MenuProps} from 'antd';
import {App, Button, Dropdown, Form, Input, Select, Space, Tag} from 'antd';
import type {ColumnsType} from 'antd/es/table';
import {Edit, Eye, EyeOff, FileWarning, GripVertical, Lock, MoreVertical, Plus, RefreshCw, Shield, Tags, Trash2} from 'lucide-react';
import {Edit, Eye, EyeOff, FileWarning, GripVertical, Lock, MoreVertical, Play, Plus, PowerOff, RefreshCw, Shield, Tags, Trash2} from 'lucide-react';
import {useMutation, useQuery, useQueryClient} from '@tanstack/react-query';
import type {DragEndEvent} from '@dnd-kit/core';
import {DndContext, KeyboardSensor, PointerSensor, useSensor, useSensors} from '@dnd-kit/core';
@@ -18,7 +18,7 @@ import {
} from '@dnd-kit/sortable';
import {CSS} from '@dnd-kit/utilities';
import dayjs from 'dayjs';
import {deleteAgent, getTags, listAgentsByAdmin, updateAgentOrder} from '@/api/agent.ts';
import {deleteAgent, getTags, listAgentsByAdmin, updateAgentEnabled, updateAgentOrder} from '@/api/agent.ts';
import type {Agent} from '@/types';
import {getErrorMessage} from '@/lib/utils';
import {PageHeader} from '@/components/PageHeader';
@@ -146,6 +146,17 @@ const AgentList = () => {
},
});
const enabledMutation = useMutation({
mutationFn: ({agentId, enabled}: {agentId: string; enabled: boolean}) => updateAgentEnabled(agentId, enabled),
onSuccess: (_response, variables) => {
messageApi.success(variables.enabled ? '主机已启用' : '主机已禁用');
queryClient.invalidateQueries({queryKey: ['admin', 'agents']});
},
onError: (error: unknown, variables) => {
messageApi.error(getErrorMessage(error, variables.enabled ? '启用主机失败' : '禁用主机失败'));
},
});
const orderMutation = useMutation({
mutationFn: (agentIds: string[]) => updateAgentOrder(agentIds),
onMutate: async (agentIds) => {
@@ -219,6 +230,23 @@ const AgentList = () => {
});
};
const handleEnabledChange = (agent: Agent) => {
if (!agent.enabled) {
enabledMutation.mutate({agentId: agent.id, enabled: true});
return;
}
modal.confirm({
title: '禁用主机',
content: `确定要禁用主机「${agent.name || agent.hostname}」吗？禁用期间接收到的数据将被忽略，也不会触发告警。`,
okText: '确认禁用',
cancelText: '取消',
okButtonProps: {danger: true},
centered: true,
onOk: () => enabledMutation.mutateAsync({agentId: agent.id, enabled: false}),
});
};
const handleBatchTags = () => {
if (selectedRowKeys.length === 0) {
messageApi.warning('请先选择要操作的探针');
@@ -270,9 +298,11 @@ const AgentList = () => {
}
// 状态过滤
if (status) {
if (status === 'disabled') {
result = result.filter((agent: Agent) => !agent.enabled);
} else if (status) {
const statusValue = status === 'online' ? 1 : 0;
result = result.filter((agent: Agent) => agent.status === statusValue);
result = result.filter((agent: Agent) => agent.enabled && agent.status === statusValue);
}
// 标签过滤
@@ -356,11 +386,16 @@ const AgentList = () => {
dataIndex: 'status',
key: 'status',
width: 80,
render: (_, record) => (
render: (_, record) => {
if (!record.enabled) {
return <Tag color="warning">已禁用</Tag>;
}
return (
<Tag color={record.status === 1 ? 'success' : 'default'}>
{record.status === 1 ? '在线' : '离线'}
</Tag>
),
);
},
},
{
title: '可见性',
@@ -554,6 +589,17 @@ const AgentList = () => {
{
type: 'divider',
},
{
key: 'enabled',
label: record.enabled ? '禁用主机' : '启用主机',
icon: record.enabled ? <PowerOff size={14}/> : <Play size={14}/>,
danger: record.enabled,
disabled: enabledMutation.isPending,
onClick: () => handleEnabledChange(record),
},
{
type: 'divider',
},
{
key: 'delete',
label: '删除探针',
@@ -621,6 +667,7 @@ const AgentList = () => {
options={[
{label: '在线', value: 'online'},
{label: '离线', value: 'offline'},
{label: '已禁用', value: 'disabled'},
]}
/>
</Form.Item>
web/src/types/index.ts
+1
View File
@@ -26,6 +26,7 @@ export interface Agent {
version: string;
tags?: string[]; // 标签
expireTime?: number; // 到期时间（时间戳毫秒）
enabled: boolean; // 是否接收并处理探针数据
status: number;
visibility?: string; // 可见性: public-匿名可见, private-登录可见
weight?: number; // 展示顺序（由拖动排序维护）