pomerium/internal/identity/manager/schedulers.go

package manager

import (
	"context"
	"sync/atomic"
	"time"

	"github.com/pomerium/pomerium/pkg/grpc/session"
)

type updateUserInfoScheduler struct {
	mgr    *Manager
	userID string
	reset  chan struct{}
	cancel context.CancelFunc
}

func newUpdateUserInfoScheduler(ctx context.Context, mgr *Manager, userID string) *updateUserInfoScheduler {
	uuis := &updateUserInfoScheduler{
		mgr:    mgr,
		userID: userID,
		reset:  make(chan struct{}, 1),
	}
	ctx = context.WithoutCancel(ctx)
	ctx, uuis.cancel = context.WithCancel(ctx)
	go uuis.run(ctx)
	return uuis
}

func (uuis *updateUserInfoScheduler) Reset() {
	select {
	case uuis.reset <- struct{}{}:
	default:
	}
}

func (uuis *updateUserInfoScheduler) Stop() {
	uuis.cancel()
}

func (uuis *updateUserInfoScheduler) run(ctx context.Context) {
	ticker := time.NewTicker(uuis.mgr.cfg.Load().updateUserInfoInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-uuis.reset:
			ticker.Reset(uuis.mgr.cfg.Load().updateUserInfoInterval)
		case <-ticker.C:
			uuis.mgr.updateUserInfo(ctx, uuis.userID)
		}
	}
}

type refreshSessionScheduler struct {
	mgr         *Manager
	sessionID   string
	lastRefresh atomic.Pointer[time.Time]
	next        chan time.Time
	cancel      context.CancelFunc
}

func newRefreshSessionScheduler(
	ctx context.Context,
	mgr *Manager,
	sessionID string,
) *refreshSessionScheduler {
	rss := &refreshSessionScheduler{
		mgr:       mgr,
		sessionID: sessionID,
		next:      make(chan time.Time, 1),
	}
	now := rss.mgr.cfg.Load().now()
	rss.lastRefresh.Store(&now)
	ctx = context.WithoutCancel(ctx)
	ctx, rss.cancel = context.WithCancel(ctx)
	go rss.run(ctx)
	return rss
}

func (rss *refreshSessionScheduler) Update(s *session.Session) {
	due := nextSessionRefresh(
		s,
		*rss.lastRefresh.Load(),
		rss.mgr.cfg.Load().sessionRefreshGracePeriod,
		rss.mgr.cfg.Load().sessionRefreshCoolOffDuration,
	)
	for {
		select {
		case <-rss.next:
		default:
		}
		select {
		case rss.next <- due:
			return
		default:
		}
	}
}

func (rss *refreshSessionScheduler) Stop() {
	rss.cancel()
}

func (rss *refreshSessionScheduler) run(ctx context.Context) {
	var timer *time.Timer

	// wait for the first update
	select {
	case <-ctx.Done():
		return
	case due := <-rss.next:
		delay := max(time.Until(due), 0)
		timer = time.NewTimer(delay)
		defer timer.Stop()
	}

	// wait for updates or for the timer to trigger
	for {
		select {
		case <-ctx.Done():
			return
		case due := <-rss.next:
			delay := max(time.Until(due), 0)
			// stop the current timer and reset it
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(delay)
		case <-timer.C:
			now := rss.mgr.cfg.Load().now()
			rss.lastRefresh.Store(&now)
			rss.mgr.refreshSession(ctx, rss.sessionID)
		}
	}
}