OpenList/drivers/aliyundrive_open/callback.go

package aliyundrive_open

import (
	"context"
	"fmt"
	"io"
	"math/rand/v2"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/OpenListTeam/OpenList/v4/internal/conf"
	"github.com/OpenListTeam/OpenList/v4/internal/errs"
	anet "github.com/OpenListTeam/OpenList/v4/internal/net"
	"github.com/OpenListTeam/OpenList/v4/internal/stream"
	"github.com/OpenListTeam/OpenList/v4/pkg/http_range"
)

const (
	defaultCallbackConcurrency = 1
	callbackAcquireTimeout     = time.Second
	callbackRequestAttempts    = 3
	callbackRetryBaseDelay     = 200 * time.Millisecond
	callbackErrorBodyLimit     = 64 << 10
)

var callbackLimiters = struct {
	sync.Mutex
	byUser map[string]*callbackLimiter
}{byUser: make(map[string]*callbackLimiter)}

type callbackLimiter struct {
	userID        string
	mu            sync.Mutex
	active        int
	nextID        uint64
	registrations map[uint64]int
	changed       chan struct{}
}

type callbackRegistration struct {
	limiter *callbackLimiter
	id      uint64
	once    sync.Once
}

type callbackPermit struct {
	limiter *callbackLimiter
	once    sync.Once
}

func normalizeCallbackConcurrency(limit int) int {
	if limit <= 0 {
		return defaultCallbackConcurrency
	}
	return limit
}

func registerCallbackLimiter(userID string, limit int) *callbackRegistration {
	callbackLimiters.Lock()
	defer callbackLimiters.Unlock()

	limiter := callbackLimiters.byUser[userID]
	if limiter == nil {
		limiter = &callbackLimiter{
			userID:        userID,
			registrations: make(map[uint64]int),
			changed:       make(chan struct{}),
		}
		callbackLimiters.byUser[userID] = limiter
	}
	limiter.mu.Lock()
	limiter.nextID++
	id := limiter.nextID
	limiter.registrations[id] = normalizeCallbackConcurrency(limit)
	limiter.signalLocked()
	limiter.mu.Unlock()
	return &callbackRegistration{limiter: limiter, id: id}
}

func (r *callbackRegistration) unregister() {
	if r == nil || r.limiter == nil {
		return
	}
	r.once.Do(func() {
		callbackLimiters.Lock()
		defer callbackLimiters.Unlock()
		r.limiter.mu.Lock()
		delete(r.limiter.registrations, r.id)
		r.limiter.signalLocked()
		if len(r.limiter.registrations) == 0 && r.limiter.active == 0 {
			delete(callbackLimiters.byUser, r.limiter.userID)
		}
		r.limiter.mu.Unlock()
	})
}

func (r *callbackRegistration) acquire(ctx context.Context) (*callbackPermit, error) {
	if r == nil || r.limiter == nil {
		return nil, errs.NewErr(errs.TemporaryCapacity, "callback limiter is unavailable")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	waitCtx, cancel := context.WithTimeout(ctx, callbackAcquireTimeout)
	defer cancel()
	for {
		r.limiter.mu.Lock()
		if r.limiter.active < r.limiter.limitLocked() {
			r.limiter.active++
			r.limiter.mu.Unlock()
			return &callbackPermit{limiter: r.limiter}, nil
		}
		changed := r.limiter.changed
		r.limiter.mu.Unlock()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-waitCtx.Done():
			if err := ctx.Err(); err != nil {
				return nil, err
			}
			return nil, errs.NewErr(errs.TemporaryCapacity, "timed out waiting for callback admission")
		case <-changed:
		}
	}
}

func (l *callbackLimiter) limitLocked() int {
	limit := 0
	for _, registered := range l.registrations {
		if limit == 0 || registered < limit {
			limit = registered
		}
	}
	return limit
}

func (l *callbackLimiter) signalLocked() {
	close(l.changed)
	l.changed = make(chan struct{})
}

func (p *callbackPermit) release() {
	if p == nil || p.limiter == nil {
		return
	}
	p.once.Do(func() {
		callbackLimiters.Lock()
		defer callbackLimiters.Unlock()
		p.limiter.mu.Lock()
		p.limiter.active--
		p.limiter.signalLocked()
		if len(p.limiter.registrations) == 0 && p.limiter.active == 0 {
			delete(callbackLimiters.byUser, p.limiter.userID)
		}
		p.limiter.mu.Unlock()
	})
}

func (d *AliyundriveOpen) callbackRegistration() *callbackRegistration {
	if d.callback != nil {
		return d.callback
	}
	if d.ref != nil {
		return d.ref.callbackRegistration()
	}
	return nil
}

func (d *AliyundriveOpen) callbackRangeReader(url string, size int64) stream.RangeReaderFunc {
	return func(ctx context.Context, requested http_range.Range) (io.ReadCloser, error) {
		if requested.Length < 0 || requested.Start+requested.Length > size {
			requested.Length = size - requested.Start
		}
		for attempt := 0; attempt < callbackRequestAttempts; attempt++ {
			permit, err := d.callbackRegistration().acquire(ctx)
			if err != nil {
				return nil, err
			}
			body, retry, err := openCallbackRange(ctx, url, size, requested)
			if !retry && err == nil {
				return newCallbackBody(ctx, body, permit.release), nil
			}
			permit.release()
			if !retry {
				return nil, err
			}
			if attempt+1 == callbackRequestAttempts {
				return nil, errs.NewErr(errs.TemporaryCapacity, "Aliyun callback concurrency limit rejected %d attempts", callbackRequestAttempts)
			}
			delay := callbackRetryBaseDelay << attempt
			delay += time.Duration(rand.Int64N(int64(delay / 2)))
			timer := time.NewTimer(delay)
			select {
			case <-ctx.Done():
				timer.Stop()
				return nil, ctx.Err()
			case <-timer.C:
			}
		}
		return nil, errs.NewErr(errs.TemporaryCapacity, "callback attempts exhausted")
	}
}

func openCallbackRange(ctx context.Context, url string, size int64, requested http_range.Range) (io.ReadCloser, bool, error) {
	requestHeader, _ := ctx.Value(conf.RequestHeaderKey).(http.Header)
	header := anet.ProcessHeader(requestHeader, nil)
	header = http_range.ApplyRangeToHttpHeader(requested, header)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, false, fmt.Errorf("create Aliyun callback request: %w", err)
	}
	req.Header = header
	response, err := anet.HttpClient().Do(req)
	if err != nil {
		return nil, false, fmt.Errorf("Aliyun callback request failed: %w", err)
	}
	if response.StatusCode >= http.StatusBadRequest {
		defer response.Body.Close()
		body, readErr := io.ReadAll(io.LimitReader(response.Body, callbackErrorBodyLimit))
		if readErr != nil {
			return nil, false, fmt.Errorf("read Aliyun callback error response: %w", readErr)
		}
		if isCallbackCapacityRejection(response.StatusCode, body) {
			return nil, true, nil
		}
		message := strings.ReplaceAll(strings.TrimSpace(string(body)), url, "<redacted>")
		return nil, false, fmt.Errorf("Aliyun callback request failed: %w; response: %s", anet.HttpStatusCodeError(response.StatusCode), message)
	}
	if requested.Start == 0 && requested.Length == size || response.StatusCode == http.StatusPartialContent || callbackContentRangeStartsAt(response.Header, requested.Start) {
		return response.Body, false, nil
	}
	if response.StatusCode == http.StatusOK {
		body, rangeErr := anet.GetRangedHttpReader(response.Body, requested.Start, requested.Length)
		if rangeErr != nil {
			response.Body.Close()
			return nil, false, rangeErr
		}
		return body, false, nil
	}
	return response.Body, false, nil
}

func isCallbackCapacityRejection(status int, body []byte) bool {
	return status == http.StatusForbidden &&
		strings.Contains(string(body), "RequestDeniedByCallback") &&
		strings.Contains(string(body), "ExceedMaxConcurrency")
}

func callbackContentRangeStartsAt(header http.Header, offset int64) bool {
	start, _, err := http_range.ParseContentRange(header.Get("Content-Range"))
	return err == nil && start == offset
}

type callbackBody struct {
	body    io.ReadCloser
	release func()
	once    sync.Once
	mu      sync.Mutex
	stop    func() bool
}

func newCallbackBody(ctx context.Context, body io.ReadCloser, release func()) *callbackBody {
	b := &callbackBody{body: body, release: release}
	stop := context.AfterFunc(ctx, func() { _ = b.Close() })
	b.mu.Lock()
	b.stop = stop
	b.mu.Unlock()
	return b
}

func (b *callbackBody) Read(p []byte) (int, error) {
	n, err := b.body.Read(p)
	if err != nil {
		_ = b.Close()
	}
	return n, err
}

func (b *callbackBody) Close() error {
	var err error
	b.once.Do(func() {
		b.mu.Lock()
		stop := b.stop
		b.mu.Unlock()
		if stop != nil {
			stop()
		}
		err = b.body.Close()
		b.release()
	})
	return err
}