mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2025-09-12 05:00:05 +00:00
feat: implement periodic traffic reset job and integrate with cron scheduler
This commit is contained in:
egregors
parent
283aa528eb
commit
8484945bb2
4 changed files with 915 additions and 0 deletions
124
web/job/periodic_traffic_reset_job.go
Normal file
124
web/job/periodic_traffic_reset_job.go
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
package job
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"x-ui/database/model"
|
||||
"x-ui/logger"
|
||||
"x-ui/web/service"
|
||||
)
|
||||
|
||||
const (
|
||||
resetTypeNever = "never"
|
||||
resetTypeDaily = "daily"
|
||||
resetTypeWeekly = "weekly"
|
||||
resetTypeMonthly = "monthly"
|
||||
)
|
||||
|
||||
var validResetTypes = map[string]bool{
|
||||
resetTypeNever: true,
|
||||
resetTypeDaily: true,
|
||||
resetTypeWeekly: true,
|
||||
resetTypeMonthly: true,
|
||||
}
|
||||
|
||||
type PeriodicTrafficResetJob struct {
|
||||
inboundService service.InboundService
|
||||
lastResetTimes map[string]time.Time
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewPeriodicTrafficResetJob() *PeriodicTrafficResetJob {
|
||||
return &PeriodicTrafficResetJob{
|
||||
lastResetTimes: make(map[string]time.Time),
|
||||
mu: sync.RWMutex{},
|
||||
}
|
||||
}
|
||||
|
||||
func (j *PeriodicTrafficResetJob) Run() {
|
||||
inbounds, err := j.inboundService.GetAllInbounds()
|
||||
if err != nil {
|
||||
logger.Warning("Failed to get inbounds for traffic reset:", err)
|
||||
return
|
||||
}
|
||||
|
||||
resetCount := 0
|
||||
now := time.Now()
|
||||
|
||||
for _, inbound := range inbounds {
|
||||
if !j.shouldResetTraffic(inbound, now) {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := j.inboundService.ResetAllClientTraffics(inbound.Id); err != nil {
|
||||
logger.Warning("Failed to reset traffic for inbound", inbound.Id, ":", err)
|
||||
continue
|
||||
}
|
||||
|
||||
j.updateLastResetTime(inbound, now)
|
||||
resetCount++
|
||||
logger.Infof("Reset traffic for inbound %d (%s)", inbound.Id, inbound.Remark)
|
||||
}
|
||||
|
||||
if resetCount > 0 {
|
||||
logger.Infof("Periodic traffic reset completed: %d inbounds reset", 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 {
|
||||
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
|
||||
}
|
||||
333
web/job/periodic_traffic_reset_job_integration_test.go
Normal file
333
web/job/periodic_traffic_reset_job_integration_test.go
Normal file
|
|
@ -0,0 +1,333 @@
|
|||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
455
web/job/periodic_traffic_reset_job_test.go
Normal file
455
web/job/periodic_traffic_reset_job_test.go
Normal file
|
|
@ -0,0 +1,455 @@
|
|||
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")
|
||||
}
|
||||
}
|
||||
|
|
@ -266,6 +266,9 @@ func (s *Server) startTask() {
|
|||
// check client ips from log file every day
|
||||
s.cron.AddJob("@daily", job.NewClearLogsJob())
|
||||
|
||||
// Check for periodic traffic resets every 10 minutes
|
||||
s.cron.AddJob("@every 10m", job.NewPeriodicTrafficResetJob())
|
||||
|
||||
// Make a traffic condition every day, 8:30
|
||||
var entry cron.EntryID
|
||||
isTgbotenabled, err := s.settingService.GetTgbotEnabled()
|
||||
|
|
|
|||
Loading…
Reference in a new issue