package ratelimited

// Code generated by gowrap. DO NOT EDIT.
// template: ../templates/ratelimited.tmpl
// gowrap: http://github.com/hexdigest/gowrap

import (
	"context"

	"github.com/uber/cadence/common/persistence"
	"github.com/uber/cadence/common/quotas"
)

// ratelimitedShardManager implements persistence.ShardManager interface instrumented with rate limiter.
type ratelimitedShardManager struct {
	wrapped     persistence.ShardManager
	rateLimiter quotas.Limiter
}

// NewShardManager creates a new instance of ShardManager with ratelimiter.
func NewShardManager(
	wrapped persistence.ShardManager,
	rateLimiter quotas.Limiter,
) persistence.ShardManager {
	return &ratelimitedShardManager{
		wrapped:     wrapped,
		rateLimiter: rateLimiter,
	}
}

func (c *ratelimitedShardManager) Close() {
	c.wrapped.Close()
	return
}

func (c *ratelimitedShardManager) CreateShard(ctx context.Context, request *persistence.CreateShardRequest) (err error) {
	if ok := c.rateLimiter.Allow(); !ok {
		err = ErrPersistenceLimitExceeded
		return
	}
	return c.wrapped.CreateShard(ctx, request)
}

func (c *ratelimitedShardManager) GetName() (s1 string) {
	return c.wrapped.GetName()
}

func (c *ratelimitedShardManager) GetShard(ctx context.Context, request *persistence.GetShardRequest) (gp1 *persistence.GetShardResponse, err error) {
	if ok := c.rateLimiter.Allow(); !ok {
		err = ErrPersistenceLimitExceeded
		return
	}
	return c.wrapped.GetShard(ctx, request)
}

func (c *ratelimitedShardManager) UpdateShard(ctx context.Context, request *persistence.UpdateShardRequest) (err error) {
	if ok := c.rateLimiter.Allow(); !ok {
		err = ErrPersistenceLimitExceeded
		return
	}
	return c.wrapped.UpdateShard(ctx, request)
}
