feat: enhance periodic traffic reset functionality with scheduling and inbound filtering

This commit is contained in:
egregors 2025-09-06 10:40:59 +02:00
parent 8484945bb2
commit 4ff0d56e8a
No known key found for this signature in database
GPG key ID: 78C8DDCC8BB0494A
5 changed files with 29 additions and 872 deletions

View file

@ -1,124 +1,49 @@
package job package job
import ( import (
"sync"
"time"
"x-ui/database/model" "x-ui/database/model"
"x-ui/logger" "x-ui/logger"
"x-ui/web/service" "x-ui/web/service"
) )
const ( type Period string
resetTypeNever = "never"
resetTypeDaily = "daily"
resetTypeWeekly = "weekly"
resetTypeMonthly = "monthly"
)
var validResetTypes = map[string]bool{
resetTypeNever: true,
resetTypeDaily: true,
resetTypeWeekly: true,
resetTypeMonthly: true,
}
type PeriodicTrafficResetJob struct { type PeriodicTrafficResetJob struct {
inboundService service.InboundService inboundService service.InboundService
lastResetTimes map[string]time.Time period Period
mu sync.RWMutex
} }
func NewPeriodicTrafficResetJob() *PeriodicTrafficResetJob { func NewPeriodicTrafficResetJob(period Period) *PeriodicTrafficResetJob {
return &PeriodicTrafficResetJob{ return &PeriodicTrafficResetJob{
lastResetTimes: make(map[string]time.Time), period: period,
mu: sync.RWMutex{},
} }
} }
func (j *PeriodicTrafficResetJob) Run() { func (j *PeriodicTrafficResetJob) Run() {
inbounds, err := j.inboundService.GetAllInbounds() inbounds, err := j.inboundService.GetInboundsByPeriodicTrafficReset(string(j.period))
logger.Infof("Running periodic traffic reset job for period: %s", j.period)
if err != nil { if err != nil {
logger.Warning("Failed to get inbounds for traffic reset:", err) logger.Warning("Failed to get inbounds for traffic reset:", err)
return return
} }
resetCount := 0 resetCount := 0
now := time.Now()
for _, inbound := range inbounds { for _, inbound := range inbounds {
if !j.shouldResetTraffic(inbound, now) {
continue
}
if err := j.inboundService.ResetAllClientTraffics(inbound.Id); err != nil { if err := j.inboundService.ResetAllClientTraffics(inbound.Id); err != nil {
logger.Warning("Failed to reset traffic for inbound", inbound.Id, ":", err) logger.Warning("Failed to reset traffic for inbound", inbound.Id, ":", err)
continue continue
} }
j.updateLastResetTime(inbound, now)
resetCount++ resetCount++
logger.Infof("Reset traffic for inbound %d (%s)", inbound.Id, inbound.Remark) logger.Infof("Reset traffic for inbound %d (%s)", inbound.Id, inbound.Remark)
} }
if resetCount > 0 { if resetCount > 0 {
logger.Infof("Periodic traffic reset completed: %d inbounds reset", resetCount) logger.Infof("Periodic traffic reset completed: %d inbounds reseted", resetCount)
} }
} }
func (j *PeriodicTrafficResetJob) shouldResetTraffic(inbound *model.Inbound, now time.Time) bool {
if !validResetTypes[inbound.PeriodicTrafficReset] || inbound.PeriodicTrafficReset == resetTypeNever {
return false
}
resetKey := j.getResetKey(inbound)
lastReset := j.getLastResetTime(resetKey)
switch inbound.PeriodicTrafficReset {
case resetTypeDaily:
return j.shouldResetDaily(now, lastReset)
case resetTypeWeekly:
return j.shouldResetWeekly(now, lastReset)
case resetTypeMonthly:
return j.shouldResetMonthly(now, lastReset)
default:
return false
}
}
func (j *PeriodicTrafficResetJob) shouldResetDaily(now, lastReset time.Time) bool {
if lastReset.IsZero() {
return now.Hour() == 0 && now.Minute() < 10
}
return now.Sub(lastReset) >= 24*time.Hour && now.Hour() == 0 && now.Minute() < 10
}
func (j *PeriodicTrafficResetJob) shouldResetWeekly(now, lastReset time.Time) bool {
if lastReset.IsZero() {
return now.Weekday() == time.Sunday && now.Hour() == 0 && now.Minute() < 10
}
return now.Sub(lastReset) >= 7*24*time.Hour && now.Weekday() == time.Sunday && now.Hour() == 0 && now.Minute() < 10
}
func (j *PeriodicTrafficResetJob) shouldResetMonthly(now, lastReset time.Time) bool {
if lastReset.IsZero() {
return now.Day() == 1 && now.Hour() == 0 && now.Minute() < 10
}
return now.Sub(lastReset) >= 28*24*time.Hour && now.Day() == 1 && now.Hour() == 0 && now.Minute() < 10
}
func (j *PeriodicTrafficResetJob) getResetKey(inbound *model.Inbound) string { func (j *PeriodicTrafficResetJob) getResetKey(inbound *model.Inbound) string {
return inbound.PeriodicTrafficReset + "_" + inbound.Tag return inbound.PeriodicTrafficReset + "_" + inbound.Tag
} }
func (j *PeriodicTrafficResetJob) getLastResetTime(key string) time.Time {
j.mu.RLock()
defer j.mu.RUnlock()
return j.lastResetTimes[key]
}
func (j *PeriodicTrafficResetJob) updateLastResetTime(inbound *model.Inbound, resetTime time.Time) {
j.mu.Lock()
defer j.mu.Unlock()
j.lastResetTimes[j.getResetKey(inbound)] = resetTime
}

View file

@ -1,333 +0,0 @@
package job
import (
"fmt"
"testing"
"time"
"x-ui/database/model"
)
// Tests to ensure the job handles various real-world scenarios properly
func TestPeriodicTrafficResetJobNewInstance(t *testing.T) {
job := NewPeriodicTrafficResetJob()
if job == nil {
t.Fatal("Expected job to be created, got nil")
}
if job.lastResetTimes == nil {
t.Fatal("Expected lastResetTimes map to be initialized")
}
if len(job.lastResetTimes) != 0 {
t.Error("Expected lastResetTimes map to be empty initially")
}
}
func TestPeriodicTrafficResetJobBoundaryConditions(t *testing.T) {
job := NewPeriodicTrafficResetJob()
// Test edge cases for time boundaries
testCases := []struct {
name string
resetType string
now time.Time
lastReset time.Time
expectedResult bool
}{
{
name: "daily reset at exact midnight boundary",
resetType: "daily",
now: time.Date(2024, 1, 15, 0, 0, 0, 0, time.UTC),
lastReset: time.Time{},
expectedResult: true,
},
{
name: "daily reset one second after midnight boundary",
resetType: "daily",
now: time.Date(2024, 1, 15, 0, 0, 1, 0, time.UTC),
lastReset: time.Time{},
expectedResult: true,
},
{
name: "weekly reset on Sunday at exact midnight",
resetType: "weekly",
now: time.Date(2024, 1, 14, 0, 0, 0, 0, time.UTC), // Sunday
lastReset: time.Time{},
expectedResult: true,
},
{
name: "monthly reset on 1st at exact midnight",
resetType: "monthly",
now: time.Date(2024, 2, 1, 0, 0, 0, 0, time.UTC),
lastReset: time.Time{},
expectedResult: true,
},
{
name: "daily reset at 10-minute boundary (should not reset)",
resetType: "daily",
now: time.Date(2024, 1, 15, 0, 10, 0, 0, time.UTC),
lastReset: time.Time{},
expectedResult: false,
},
{
name: "weekly reset at 10-minute boundary (should not reset)",
resetType: "weekly",
now: time.Date(2024, 1, 14, 0, 10, 0, 0, time.UTC), // Sunday
lastReset: time.Time{},
expectedResult: false,
},
{
name: "monthly reset at 10-minute boundary (should not reset)",
resetType: "monthly",
now: time.Date(2024, 2, 1, 0, 10, 0, 0, time.UTC),
lastReset: time.Time{},
expectedResult: false,
},
{
name: "leap year handling - February 29th to March 1st",
resetType: "monthly",
now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), // 29 days ago
expectedResult: true,
},
{
name: "weekly reset after 6 days and 23 hours (should not reset)",
resetType: "weekly",
now: time.Date(2024, 1, 20, 23, 0, 0, 0, time.UTC), // Saturday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday
expectedResult: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
inbound := &model.Inbound{
Id: 1,
Tag: "test-inbound",
PeriodicTrafficReset: tc.resetType,
}
// Set up last reset time if provided
if !tc.lastReset.IsZero() {
job.updateLastResetTime(inbound, tc.lastReset)
}
result := job.shouldResetTraffic(inbound, tc.now)
if result != tc.expectedResult {
t.Errorf("Expected %v, got %v", tc.expectedResult, result)
}
// Clean up for next test
job.lastResetTimes = make(map[string]time.Time)
})
}
}
func TestPeriodicTrafficResetJobTimezoneHandling(t *testing.T) {
job := NewPeriodicTrafficResetJob()
inbound := &model.Inbound{
Id: 1,
Tag: "timezone-test",
PeriodicTrafficReset: "daily",
}
// Test that UTC midnight behaves consistently
utcMidnight := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC)
result := job.shouldResetTraffic(inbound, utcMidnight)
if !result {
t.Error("Expected reset at UTC midnight")
}
// Test that the same instant in different timezones produces the same absolute behavior
// This tests consistency rather than timezone conversion
utcTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC)
utcTimestamp := utcTime.Unix()
// Create equivalent time from timestamp
equivalentTime := time.Unix(utcTimestamp, 0).UTC()
result1 := job.shouldResetTraffic(inbound, utcTime)
result2 := job.shouldResetTraffic(inbound, equivalentTime)
if result1 != result2 {
t.Errorf("Equivalent times should behave consistently: %v vs %v", result1, result2)
}
}
func TestPeriodicTrafficResetJobConcurrentSafety(t *testing.T) {
job := NewPeriodicTrafficResetJob()
inbound := &model.Inbound{
Id: 1,
Tag: "concurrent-test",
PeriodicTrafficReset: "daily",
}
// Test concurrent operations
done := make(chan bool, 100)
// Start multiple goroutines performing different operations
for i := 0; i < 50; i++ {
go func(i int) {
defer func() { done <- true }()
// Mix of operations
testTime := time.Date(2024, 1, 15+i%5, 0, 5, 0, 0, time.UTC)
job.updateLastResetTime(inbound, testTime)
_ = job.getLastResetTime(job.getResetKey(inbound))
_ = job.shouldResetTraffic(inbound, testTime)
}(i)
}
// Wait for all goroutines to complete
for i := 0; i < 50; i++ {
select {
case <-done:
// Good
case <-time.After(5 * time.Second):
t.Fatal("Concurrent test timed out - possible deadlock")
}
}
}
func TestPeriodicTrafficResetJobValidationRobustness(t *testing.T) {
job := NewPeriodicTrafficResetJob()
testCases := []struct {
name string
inbound *model.Inbound
now time.Time
expected bool
}{
{
name: "nil inbound",
inbound: nil,
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
// This should panic or handle gracefully - depends on implementation
},
{
name: "empty tag",
inbound: &model.Inbound{
Id: 1,
Tag: "",
PeriodicTrafficReset: "daily",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true, // Should still work with empty tag
},
{
name: "very long tag",
inbound: &model.Inbound{
Id: 1,
Tag: "very-very-very-long-tag-name-that-might-cause-issues-with-memory-or-key-generation-in-the-reset-tracking-system",
PeriodicTrafficReset: "daily",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "special characters in tag",
inbound: &model.Inbound{
Id: 1,
Tag: "test-inbound_with.special@chars#123",
PeriodicTrafficReset: "daily",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "zero time",
inbound: &model.Inbound{
Id: 1,
Tag: "zero-time-test",
PeriodicTrafficReset: "daily",
},
now: time.Time{}, // Zero time
expected: true, // Zero time has hour=0, minute=0, so should reset according to logic
},
{
name: "far future time",
inbound: &model.Inbound{
Id: 1,
Tag: "future-test",
PeriodicTrafficReset: "daily",
},
now: time.Date(2099, 12, 31, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "far past time",
inbound: &model.Inbound{
Id: 1,
Tag: "past-test",
PeriodicTrafficReset: "daily",
},
now: time.Date(1970, 1, 1, 0, 5, 0, 0, time.UTC),
expected: true,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
defer func() {
if r := recover(); r != nil {
if tc.inbound == nil {
// Expected panic for nil inbound
t.Log("Expected panic for nil inbound:", r)
} else {
t.Errorf("Unexpected panic: %v", r)
}
}
}()
if tc.inbound == nil {
// Skip execution for nil test case
return
}
result := job.shouldResetTraffic(tc.inbound, tc.now)
if result != tc.expected {
t.Errorf("Expected %v, got %v", tc.expected, result)
}
})
}
}
func TestPeriodicTrafficResetJobMemoryManagement(t *testing.T) {
job := NewPeriodicTrafficResetJob()
// Test with many inbounds to check memory usage
const numInbounds = 1000
testTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC)
for i := 0; i < numInbounds; i++ {
inbound := &model.Inbound{
Id: i,
Tag: fmt.Sprintf("test-inbound-%d", i),
PeriodicTrafficReset: "daily",
}
job.updateLastResetTime(inbound, testTime)
}
// Verify all entries are stored
if len(job.lastResetTimes) != numInbounds {
t.Errorf("Expected %d entries in lastResetTimes, got %d", numInbounds, len(job.lastResetTimes))
}
// Test retrieval of all entries
for i := 0; i < numInbounds; i++ {
inbound := &model.Inbound{
Id: i,
Tag: fmt.Sprintf("test-inbound-%d", i),
PeriodicTrafficReset: "daily",
}
retrievedTime := job.getLastResetTime(job.getResetKey(inbound))
if !retrievedTime.Equal(testTime) {
t.Errorf("Expected time %v for inbound %d, got %v", testTime, i, retrievedTime)
}
}
}

View file

@ -1,455 +0,0 @@
package job
import (
"testing"
"time"
"x-ui/database/model"
)
func TestShouldResetDaily(t *testing.T) {
job := NewPeriodicTrafficResetJob()
tests := []struct {
name string
now time.Time
lastReset time.Time
expected bool
}{
{
name: "first time reset at midnight",
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
lastReset: time.Time{}, // zero time
expected: true,
},
{
name: "first time reset not at midnight",
now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC),
lastReset: time.Time{}, // zero time
expected: false,
},
{
name: "reset after 24 hours at midnight",
now: time.Date(2024, 1, 16, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "reset after 24 hours but not at midnight",
now: time.Date(2024, 1, 16, 12, 0, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset same day",
now: time.Date(2024, 1, 15, 0, 8, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset after 23 hours",
now: time.Date(2024, 1, 15, 23, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "midnight but minute > 10",
now: time.Date(2024, 1, 16, 0, 15, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "edge case: exactly at minute 10",
now: time.Date(2024, 1, 16, 0, 10, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "edge case: minute 9 (last valid minute)",
now: time.Date(2024, 1, 16, 0, 9, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := job.shouldResetDaily(tt.now, tt.lastReset)
if result != tt.expected {
t.Errorf("shouldResetDaily() = %v, expected %v", result, tt.expected)
}
})
}
}
func TestShouldResetWeekly(t *testing.T) {
job := NewPeriodicTrafficResetJob()
tests := []struct {
name string
now time.Time
lastReset time.Time
expected bool
}{
{
name: "first time reset on Sunday at midnight",
now: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Sunday
lastReset: time.Time{}, // zero time
expected: true,
},
{
name: "first time reset on Monday",
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), // Monday
lastReset: time.Time{}, // zero time
expected: false,
},
{
name: "reset after 7 days on Sunday at midnight",
now: time.Date(2024, 1, 21, 0, 5, 0, 0, time.UTC), // Sunday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday
expected: true,
},
{
name: "reset after 7 days on Sunday but not at midnight",
now: time.Date(2024, 1, 21, 12, 0, 0, 0, time.UTC), // Sunday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset after 6 days on Saturday",
now: time.Date(2024, 1, 20, 0, 5, 0, 0, time.UTC), // Saturday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset same Sunday",
now: time.Date(2024, 1, 14, 0, 8, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset after 8 days but on Monday",
now: time.Date(2024, 1, 22, 0, 5, 0, 0, time.UTC), // Monday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday
expected: false,
},
{
name: "edge case: Sunday but minute > 10",
now: time.Date(2024, 1, 21, 0, 15, 0, 0, time.UTC), // Sunday
lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC),
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := job.shouldResetWeekly(tt.now, tt.lastReset)
if result != tt.expected {
t.Errorf("shouldResetWeekly() = %v, expected %v", result, tt.expected)
}
})
}
}
func TestShouldResetMonthly(t *testing.T) {
job := NewPeriodicTrafficResetJob()
tests := []struct {
name string
now time.Time
lastReset time.Time
expected bool
}{
{
name: "first time reset on 1st day at midnight",
now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC),
lastReset: time.Time{}, // zero time
expected: true,
},
{
name: "first time reset on 15th day",
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
lastReset: time.Time{}, // zero time
expected: false,
},
{
name: "reset after 28 days on 1st day at midnight",
now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC), // 28 days ago
expected: true,
},
{
name: "reset after 32 days on 1st day at midnight",
now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 29, 0, 5, 0, 0, time.UTC), // 32 days ago
expected: true,
},
{
name: "reset after 28 days on 1st day but not at midnight",
now: time.Date(2024, 2, 1, 12, 0, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset after 27 days",
now: time.Date(2024, 1, 31, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset same day (1st)",
now: time.Date(2024, 1, 1, 0, 8, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 1, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "reset after 28 days but on 2nd day",
now: time.Date(2024, 2, 2, 0, 5, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 5, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "leap year February edge case",
now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC), // March 1st in leap year
lastReset: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), // February 1st (29 days ago)
expected: true,
},
{
name: "edge case: 1st day but minute > 10",
now: time.Date(2024, 2, 1, 0, 15, 0, 0, time.UTC),
lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC),
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := job.shouldResetMonthly(tt.now, tt.lastReset)
if result != tt.expected {
t.Errorf("shouldResetMonthly() = %v, expected %v", result, tt.expected)
}
})
}
}
func TestShouldResetTraffic(t *testing.T) {
tests := []struct {
name string
inbound *model.Inbound
now time.Time
setup func(*PeriodicTrafficResetJob, *model.Inbound)
expected bool
}{
{
name: "never reset type",
inbound: &model.Inbound{
Id: 1,
Tag: "test1",
PeriodicTrafficReset: "never",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "invalid reset type",
inbound: &model.Inbound{
Id: 2,
Tag: "test2",
PeriodicTrafficReset: "invalid",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: false,
},
{
name: "daily reset should trigger",
inbound: &model.Inbound{
Id: 3,
Tag: "test3",
PeriodicTrafficReset: "daily",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "weekly reset should trigger",
inbound: &model.Inbound{
Id: 4,
Tag: "test4",
PeriodicTrafficReset: "weekly",
},
now: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Sunday
expected: true,
},
{
name: "monthly reset should trigger",
inbound: &model.Inbound{
Id: 5,
Tag: "test5",
PeriodicTrafficReset: "monthly",
},
now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC),
expected: true,
},
{
name: "daily reset with existing reset time - should not trigger",
inbound: &model.Inbound{
Id: 6,
Tag: "test6",
PeriodicTrafficReset: "daily",
},
now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC),
setup: func(j *PeriodicTrafficResetJob, inbound *model.Inbound) {
j.updateLastResetTime(inbound, time.Date(2024, 1, 15, 0, 3, 0, 0, time.UTC))
},
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Create new job for each test to avoid state interference
testJob := NewPeriodicTrafficResetJob()
if tt.setup != nil {
tt.setup(testJob, tt.inbound)
}
result := testJob.shouldResetTraffic(tt.inbound, tt.now)
if result != tt.expected {
t.Errorf("shouldResetTraffic() = %v, expected %v", result, tt.expected)
}
})
}
}
func TestGetResetKey(t *testing.T) {
job := NewPeriodicTrafficResetJob()
tests := []struct {
name string
inbound *model.Inbound
expected string
}{
{
name: "daily reset key",
inbound: &model.Inbound{
Tag: "test-inbound",
PeriodicTrafficReset: "daily",
},
expected: "daily_test-inbound",
},
{
name: "weekly reset key",
inbound: &model.Inbound{
Tag: "proxy-1",
PeriodicTrafficReset: "weekly",
},
expected: "weekly_proxy-1",
},
{
name: "monthly reset key",
inbound: &model.Inbound{
Tag: "vless-tcp",
PeriodicTrafficReset: "monthly",
},
expected: "monthly_vless-tcp",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := job.getResetKey(tt.inbound)
if result != tt.expected {
t.Errorf("getResetKey() = %v, expected %v", result, tt.expected)
}
})
}
}
func TestLastResetTimeOperations(t *testing.T) {
job := NewPeriodicTrafficResetJob()
inbound := &model.Inbound{
Id: 1,
Tag: "test-inbound",
PeriodicTrafficReset: "daily",
}
// Test getting non-existent reset time (should return zero time)
resetTime := job.getLastResetTime(job.getResetKey(inbound))
if !resetTime.IsZero() {
t.Errorf("Expected zero time for non-existent key, got %v", resetTime)
}
// Test updating and getting reset time
testTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC)
job.updateLastResetTime(inbound, testTime)
retrievedTime := job.getLastResetTime(job.getResetKey(inbound))
if !retrievedTime.Equal(testTime) {
t.Errorf("Expected %v, got %v", testTime, retrievedTime)
}
// Test updating existing reset time
newTime := time.Date(2024, 1, 16, 0, 5, 0, 0, time.UTC)
job.updateLastResetTime(inbound, newTime)
retrievedTime = job.getLastResetTime(job.getResetKey(inbound))
if !retrievedTime.Equal(newTime) {
t.Errorf("Expected %v, got %v", newTime, retrievedTime)
}
}
// Test for concurrent access to lastResetTimes map
func TestConcurrentAccess(t *testing.T) {
job := NewPeriodicTrafficResetJob()
inbound := &model.Inbound{
Id: 1,
Tag: "concurrent-test",
PeriodicTrafficReset: "daily",
}
// Run multiple goroutines to test concurrent access
done := make(chan bool, 20)
// 10 writers
for i := 0; i < 10; i++ {
go func(i int) {
testTime := time.Date(2024, 1, 15+i, 0, 5, 0, 0, time.UTC)
job.updateLastResetTime(inbound, testTime)
done <- true
}(i)
}
// 10 readers
for i := 0; i < 10; i++ {
go func() {
_ = job.getLastResetTime(job.getResetKey(inbound))
done <- true
}()
}
// Wait for all goroutines to complete
for i := 0; i < 20; i++ {
<-done
}
// If we reach here without panic/deadlock, the concurrent access is safe
t.Log("Concurrent access test passed")
}
func TestValidResetTypes(t *testing.T) {
expectedTypes := []string{"never", "daily", "weekly", "monthly"}
for _, resetType := range expectedTypes {
if !validResetTypes[resetType] {
t.Errorf("Expected reset type %s to be valid", resetType)
}
}
// Test invalid type
if validResetTypes["invalid"] {
t.Error("Expected 'invalid' to not be a valid reset type")
}
}

View file

@ -41,6 +41,16 @@ func (s *InboundService) GetAllInbounds() ([]*model.Inbound, error) {
return inbounds, nil return inbounds, nil
} }
func (s *InboundService) GetInboundsByPeriodicTrafficReset(period string) ([]*model.Inbound, error) {
db := database.GetDB()
var inbounds []*model.Inbound
err := db.Model(model.Inbound{}).Where("periodic_traffic_reset = ?", period).Find(&inbounds).Error
if err != nil && err != gorm.ErrRecordNotFound {
return nil, err
}
return inbounds, nil
}
func (s *InboundService) checkPortExist(listen string, port int, ignoreId int) (bool, error) { func (s *InboundService) checkPortExist(listen string, port int, ignoreId int) (bool, error) {
db := database.GetDB() db := database.GetDB()
if listen == "" || listen == "0.0.0.0" || listen == "::" || listen == "::0" { if listen == "" || listen == "0.0.0.0" || listen == "::" || listen == "::0" {
@ -397,6 +407,7 @@ func (s *InboundService) UpdateInbound(inbound *model.Inbound) (*model.Inbound,
oldInbound.Remark = inbound.Remark oldInbound.Remark = inbound.Remark
oldInbound.Enable = inbound.Enable oldInbound.Enable = inbound.Enable
oldInbound.ExpiryTime = inbound.ExpiryTime oldInbound.ExpiryTime = inbound.ExpiryTime
oldInbound.PeriodicTrafficReset = inbound.PeriodicTrafficReset
oldInbound.Listen = inbound.Listen oldInbound.Listen = inbound.Listen
oldInbound.Port = inbound.Port oldInbound.Port = inbound.Port
oldInbound.Protocol = inbound.Protocol oldInbound.Protocol = inbound.Protocol

View file

@ -266,8 +266,17 @@ func (s *Server) startTask() {
// check client ips from log file every day // check client ips from log file every day
s.cron.AddJob("@daily", job.NewClearLogsJob()) s.cron.AddJob("@daily", job.NewClearLogsJob())
// Check for periodic traffic resets every 10 minutes // Periodic traffic resets
s.cron.AddJob("@every 10m", job.NewPeriodicTrafficResetJob()) logger.Info("Scheduling periodic traffic reset jobs")
{
// Run once a day, midnight
// s.cron.AddJob("@daily", job.NewPeriodicTrafficResetJob("daily"))
s.cron.AddJob("* * * * *", job.NewPeriodicTrafficResetJob("daily"))
// Run once a week, midnight between Sat/Sun
s.cron.AddJob("@weekly", job.NewPeriodicTrafficResetJob("weekly"))
// Run once a month, midnight, first of month
s.cron.AddJob("@monthly", job.NewPeriodicTrafficResetJob("monthly"))
}
// Make a traffic condition every day, 8:30 // Make a traffic condition every day, 8:30
var entry cron.EntryID var entry cron.EntryID