mirror of
https://github.com/pomerium/pomerium.git
synced 2025-04-30 19:06:33 +02:00
fe31799eb5
This also replaces instances where we manually write "return ctx.Err()" with "return context.Cause(ctx)" which is functionally identical, but will also correctly propagate cause errors if present.
60 lines
1.4 KiB
Go
60 lines
1.4 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"github.com/pomerium/pomerium/pkg/grpc/databroker"
|
|
)
|
|
|
|
type leaser struct {
|
|
client databroker.DataBrokerServiceClient
|
|
funcs []func(context.Context, databroker.DataBrokerServiceClient) error
|
|
}
|
|
|
|
// GetDataBrokerServiceClient implements the databroker.LeaseHandler interface.
|
|
func (c *leaser) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {
|
|
return c.client
|
|
}
|
|
|
|
// RunLeased implements the databroker.LeaseHandler interface.
|
|
func (c *leaser) RunLeased(ctx context.Context) error {
|
|
eg, ctx := errgroup.WithContext(ctx)
|
|
for _, fn := range c.funcs {
|
|
fn := fn
|
|
eg.Go(func() error { return fn(ctx, c.client) })
|
|
}
|
|
err := eg.Wait()
|
|
return err
|
|
}
|
|
|
|
func WithLease(
|
|
funcs ...func(context.Context, databroker.DataBrokerServiceClient) error,
|
|
) func(context.Context, databroker.DataBrokerServiceClient) error {
|
|
return func(ctx context.Context, client databroker.DataBrokerServiceClient) error {
|
|
srv := &leaser{
|
|
client: client,
|
|
funcs: funcs,
|
|
}
|
|
leaser := databroker.NewLeaser("zero-ctrl", time.Second*30, srv)
|
|
return leaser.Run(ctx)
|
|
}
|
|
}
|
|
|
|
type LeaseStatus struct {
|
|
v atomic.Bool
|
|
}
|
|
|
|
func (w *LeaseStatus) HasLease() bool {
|
|
return w.v.Load()
|
|
}
|
|
|
|
func (w *LeaseStatus) MonitorLease(ctx context.Context, _ databroker.DataBrokerServiceClient) error {
|
|
w.v.Store(true)
|
|
<-ctx.Done()
|
|
w.v.Store(false)
|
|
return context.Cause(ctx)
|
|
}
|