2023-02-15 01:52:11 +00:00
|
|
|
package logstore
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
|
|
|
|
"github.com/zitadel/logging"
|
|
|
|
|
|
|
|
"github.com/zitadel/zitadel/internal/repository/quota"
|
|
|
|
)
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
type UsageStorer[T LogRecord[T]] interface {
|
|
|
|
LogEmitter[T]
|
2023-02-15 01:52:11 +00:00
|
|
|
QuotaUnit() quota.Unit
|
|
|
|
}
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
type Service[T LogRecord[T]] struct {
|
|
|
|
queries Queries
|
|
|
|
usageStorer UsageStorer[T]
|
|
|
|
enabledSinks []*emitter[T]
|
2023-02-15 01:52:11 +00:00
|
|
|
sinkEnabled bool
|
|
|
|
reportingEnabled bool
|
|
|
|
}
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
type Queries interface {
|
|
|
|
GetRemainingQuotaUsage(ctx context.Context, instanceID string, unit quota.Unit) (remaining *uint64, err error)
|
|
|
|
}
|
|
|
|
|
|
|
|
func New[T LogRecord[T]](queries Queries, usageQuerierSink *emitter[T], additionalSink ...*emitter[T]) *Service[T] {
|
|
|
|
var usageStorer UsageStorer[T]
|
2023-02-15 01:52:11 +00:00
|
|
|
if usageQuerierSink != nil {
|
2023-09-15 14:58:45 +00:00
|
|
|
usageStorer = usageQuerierSink.emitter.(UsageStorer[T])
|
2023-02-15 01:52:11 +00:00
|
|
|
}
|
2023-09-15 14:58:45 +00:00
|
|
|
svc := &Service[T]{
|
|
|
|
queries: queries,
|
2023-02-15 01:52:11 +00:00
|
|
|
reportingEnabled: usageQuerierSink != nil && usageQuerierSink.enabled,
|
2023-09-15 14:58:45 +00:00
|
|
|
usageStorer: usageStorer,
|
2023-02-15 01:52:11 +00:00
|
|
|
}
|
2023-09-15 14:58:45 +00:00
|
|
|
for _, s := range append([]*emitter[T]{usageQuerierSink}, additionalSink...) {
|
2023-02-15 01:52:11 +00:00
|
|
|
if s != nil && s.enabled {
|
|
|
|
svc.enabledSinks = append(svc.enabledSinks, s)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
svc.sinkEnabled = len(svc.enabledSinks) > 0
|
|
|
|
return svc
|
|
|
|
}
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
func (s *Service[T]) Enabled() bool {
|
2023-02-15 01:52:11 +00:00
|
|
|
return s.sinkEnabled
|
|
|
|
}
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
func (s *Service[T]) Handle(ctx context.Context, record T) {
|
2023-02-15 01:52:11 +00:00
|
|
|
for _, sink := range s.enabledSinks {
|
|
|
|
logging.OnError(sink.Emit(ctx, record.Normalize())).WithField("record", record).Warn("failed to emit log record")
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-09-15 14:58:45 +00:00
|
|
|
func (s *Service[T]) Limit(ctx context.Context, instanceID string) *uint64 {
|
2023-02-15 01:52:11 +00:00
|
|
|
var err error
|
|
|
|
defer func() {
|
2023-09-15 14:58:45 +00:00
|
|
|
logging.OnError(err).Warn("failed to check if usage should be limited")
|
2023-02-15 01:52:11 +00:00
|
|
|
}()
|
|
|
|
if !s.reportingEnabled || instanceID == "" {
|
|
|
|
return nil
|
|
|
|
}
|
2023-09-15 14:58:45 +00:00
|
|
|
remaining, err := s.queries.GetRemainingQuotaUsage(ctx, instanceID, s.usageStorer.QuotaUnit())
|
2023-02-15 01:52:11 +00:00
|
|
|
if err != nil {
|
2023-09-15 14:58:45 +00:00
|
|
|
// TODO: shouldn't we just limit then or return the error and decide there?
|
2023-02-15 01:52:11 +00:00
|
|
|
return nil
|
|
|
|
}
|
2023-02-27 21:36:43 +00:00
|
|
|
return remaining
|
|
|
|
}
|