feat: publish sanitized custom.97 source

service/log_export.go

@@ -0,0 +1,66 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+)
+
+const logExportLockTTL = time.Minute
+
+const logExportLockRenewScript = `
+if redis.call('GET', KEYS[1]) == ARGV[1] then
+  return redis.call('PEXPIRE', KEYS[1], ARGV[2])
+end
+return 0`
+
+const logExportLockReleaseScript = `
+if redis.call('GET', KEYS[1]) == ARGV[1] then
+  return redis.call('DEL', KEYS[1])
+end
+return 0`
+
+// logExportLocks holds the running exports when Redis is disabled.
+var logExportLocks sync.Map
+
+// AcquireLogExportLock reserves the one log export a user may run at a time.
+// With Redis the reservation spans all nodes and lapses within a minute if
+// this node dies. ok is false while another export of the user is running;
+// otherwise release must be called when the export ends.
+func AcquireLogExportLock(userId int) (release func(), ok bool, err error) {
+	if !common.RedisEnabled {
+		if _, running := logExportLocks.LoadOrStore(userId, struct{}{}); running {
+			return nil, false, nil
+		}
+		return func() { logExportLocks.Delete(userId) }, true, nil
+	}
+
+	ctx := context.Background()
+	key := fmt.Sprintf("log_export:%d", userId)
+	owner := common.GetRandomString(16)
+	acquired, err := common.RDB.SetNX(ctx, key, owner, logExportLockTTL).Result()
+	if err != nil || !acquired {
+		return nil, false, err
+	}
+
+	stop := make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(logExportLockTTL / 3)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-stop:
+				return
+			case <-ticker.C:
+				common.RDB.Eval(ctx, logExportLockRenewScript, []string{key}, owner, logExportLockTTL.Milliseconds())
+			}
+		}
+	}()
+	return func() {
+		close(stop)
+		common.RDB.Eval(ctx, logExportLockReleaseScript, []string{key}, owner)
+	}, true, nil
+}