mirror of
https://github.com/pomerium/pomerium.git
synced 2025-08-03 00:40:25 +02:00
storage: add sync querier (#5570)
* storage: add fallback querier * storage: add sync querier * storage: add typed querier * use synced querier
This commit is contained in:
parent
e1d84a1dde
commit
8738066ce4
19 changed files with 569 additions and 214 deletions
184
pkg/storage/querier_sync.go
Normal file
184
pkg/storage/querier_sync.go
Normal file
|
@ -0,0 +1,184 @@
|
|||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
grpc "google.golang.org/grpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
status "google.golang.org/grpc/status"
|
||||
|
||||
"github.com/pomerium/pomerium/internal/log"
|
||||
"github.com/pomerium/pomerium/pkg/grpc/databroker"
|
||||
)
|
||||
|
||||
type syncQuerier struct {
|
||||
client databroker.DataBrokerServiceClient
|
||||
recordType string
|
||||
|
||||
cancel context.CancelFunc
|
||||
serverVersion uint64
|
||||
latestRecordVersion uint64
|
||||
|
||||
mu sync.RWMutex
|
||||
ready bool
|
||||
records RecordCollection
|
||||
}
|
||||
|
||||
// NewSyncQuerier creates a new Querier backed by an in-memory record collection
|
||||
// filled via sync calls to the databroker.
|
||||
func NewSyncQuerier(
|
||||
client databroker.DataBrokerServiceClient,
|
||||
recordType string,
|
||||
) Querier {
|
||||
q := &syncQuerier{
|
||||
client: client,
|
||||
recordType: recordType,
|
||||
records: NewRecordCollection(),
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
q.cancel = cancel
|
||||
go q.run(ctx)
|
||||
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *syncQuerier) InvalidateCache(_ context.Context, _ *databroker.QueryRequest) {
|
||||
// do nothing
|
||||
}
|
||||
|
||||
func (q *syncQuerier) Query(_ context.Context, req *databroker.QueryRequest, _ ...grpc.CallOption) (*databroker.QueryResponse, error) {
|
||||
q.mu.RLock()
|
||||
if !q.canHandleQueryLocked(req) {
|
||||
q.mu.RUnlock()
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
defer q.mu.RUnlock()
|
||||
return QueryRecordCollections(map[string]RecordCollection{
|
||||
q.recordType: q.records,
|
||||
}, req)
|
||||
}
|
||||
|
||||
func (q *syncQuerier) Stop() {
|
||||
q.cancel()
|
||||
}
|
||||
|
||||
func (q *syncQuerier) canHandleQueryLocked(req *databroker.QueryRequest) bool {
|
||||
if !q.ready {
|
||||
return false
|
||||
}
|
||||
if req.GetType() != q.recordType {
|
||||
return false
|
||||
}
|
||||
if req.MinimumRecordVersionHint != nil && q.latestRecordVersion < *req.MinimumRecordVersionHint {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (q *syncQuerier) run(ctx context.Context) {
|
||||
bo := backoff.WithContext(backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(0)), ctx)
|
||||
_ = backoff.RetryNotify(func() error {
|
||||
if q.serverVersion == 0 {
|
||||
err := q.syncLatest(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return q.sync(ctx)
|
||||
}, bo, func(err error, d time.Duration) {
|
||||
log.Ctx(ctx).Error().
|
||||
Err(err).
|
||||
Dur("delay", d).
|
||||
Msg("storage/sync-querier: error syncing records")
|
||||
})
|
||||
}
|
||||
|
||||
func (q *syncQuerier) syncLatest(ctx context.Context) error {
|
||||
stream, err := q.client.SyncLatest(ctx, &databroker.SyncLatestRequest{
|
||||
Type: q.recordType,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("error starting sync latest stream: %w", err)
|
||||
}
|
||||
|
||||
q.mu.Lock()
|
||||
q.ready = false
|
||||
q.records.Clear()
|
||||
q.mu.Unlock()
|
||||
|
||||
for {
|
||||
res, err := stream.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
} else if err != nil {
|
||||
return fmt.Errorf("error receiving sync latest message: %w", err)
|
||||
}
|
||||
|
||||
switch res := res.Response.(type) {
|
||||
case *databroker.SyncLatestResponse_Record:
|
||||
q.mu.Lock()
|
||||
q.records.Put(res.Record)
|
||||
q.mu.Unlock()
|
||||
case *databroker.SyncLatestResponse_Versions:
|
||||
q.mu.Lock()
|
||||
q.serverVersion = res.Versions.ServerVersion
|
||||
q.latestRecordVersion = res.Versions.LatestRecordVersion
|
||||
q.mu.Unlock()
|
||||
default:
|
||||
return fmt.Errorf("unknown message type from sync latest: %T", res)
|
||||
}
|
||||
}
|
||||
|
||||
q.mu.Lock()
|
||||
log.Ctx(ctx).Info().
|
||||
Str("record-type", q.recordType).
|
||||
Int("record-count", q.records.Len()).
|
||||
Uint64("latest-record-version", q.latestRecordVersion).
|
||||
Msg("storage/sync-querier: synced latest records")
|
||||
q.ready = true
|
||||
q.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (q *syncQuerier) sync(ctx context.Context) error {
|
||||
q.mu.RLock()
|
||||
req := &databroker.SyncRequest{
|
||||
ServerVersion: q.serverVersion,
|
||||
RecordVersion: q.latestRecordVersion,
|
||||
Type: q.recordType,
|
||||
}
|
||||
q.mu.RUnlock()
|
||||
|
||||
stream, err := q.client.Sync(ctx, req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error starting sync stream: %w", err)
|
||||
}
|
||||
|
||||
for {
|
||||
res, err := stream.Recv()
|
||||
if status.Code(err) == codes.Aborted {
|
||||
// this indicates the server version changed, so we need to reset
|
||||
q.mu.Lock()
|
||||
q.serverVersion = 0
|
||||
q.latestRecordVersion = 0
|
||||
q.mu.Unlock()
|
||||
return fmt.Errorf("stream was aborted due to mismatched server versions: %w", err)
|
||||
} else if err != nil {
|
||||
return fmt.Errorf("error receiving sync message: %w", err)
|
||||
}
|
||||
|
||||
q.mu.Lock()
|
||||
q.latestRecordVersion = max(q.latestRecordVersion, res.Record.Version)
|
||||
q.records.Put(res.Record)
|
||||
q.mu.Unlock()
|
||||
}
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue