1// Copyright 2016 The etcd Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package v3rpc
16
17import (
18	"github.com/coreos/etcd/etcdserver"
19	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
20	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
21	"github.com/coreos/etcd/pkg/types"
22	"golang.org/x/net/context"
23)
24
25type quotaKVServer struct {
26	pb.KVServer
27	qa quotaAlarmer
28}
29
30type quotaAlarmer struct {
31	q  etcdserver.Quota
32	a  Alarmer
33	id types.ID
34}
35
36// check whether request satisfies the quota. If there is not enough space,
37// ignore request and raise the free space alarm.
38func (qa *quotaAlarmer) check(ctx context.Context, r interface{}) error {
39	if qa.q.Available(r) {
40		return nil
41	}
42	req := &pb.AlarmRequest{
43		MemberID: uint64(qa.id),
44		Action:   pb.AlarmRequest_ACTIVATE,
45		Alarm:    pb.AlarmType_NOSPACE,
46	}
47	qa.a.Alarm(ctx, req)
48	return rpctypes.ErrGRPCNoSpace
49}
50
51func NewQuotaKVServer(s *etcdserver.EtcdServer) pb.KVServer {
52	return &quotaKVServer{
53		NewKVServer(s),
54		quotaAlarmer{etcdserver.NewBackendQuota(s), s, s.ID()},
55	}
56}
57
58func (s *quotaKVServer) Put(ctx context.Context, r *pb.PutRequest) (*pb.PutResponse, error) {
59	if err := s.qa.check(ctx, r); err != nil {
60		return nil, err
61	}
62	return s.KVServer.Put(ctx, r)
63}
64
65func (s *quotaKVServer) Txn(ctx context.Context, r *pb.TxnRequest) (*pb.TxnResponse, error) {
66	if err := s.qa.check(ctx, r); err != nil {
67		return nil, err
68	}
69	return s.KVServer.Txn(ctx, r)
70}
71
72type quotaLeaseServer struct {
73	pb.LeaseServer
74	qa quotaAlarmer
75}
76
77func (s *quotaLeaseServer) LeaseGrant(ctx context.Context, cr *pb.LeaseGrantRequest) (*pb.LeaseGrantResponse, error) {
78	if err := s.qa.check(ctx, cr); err != nil {
79		return nil, err
80	}
81	return s.LeaseServer.LeaseGrant(ctx, cr)
82}
83
84func NewQuotaLeaseServer(s *etcdserver.EtcdServer) pb.LeaseServer {
85	return &quotaLeaseServer{
86		NewLeaseServer(s),
87		quotaAlarmer{etcdserver.NewBackendQuota(s), s, s.ID()},
88	}
89}
90