blob: a145b8b0911e6fa7cda0c3414c6a736a35c07d49 [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// 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 "context"
19
20 "go.etcd.io/etcd/etcdserver"
21 "go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
22 pb "go.etcd.io/etcd/etcdserver/etcdserverpb"
23 "go.etcd.io/etcd/pkg/types"
24)
25
26type quotaKVServer struct {
27 pb.KVServer
28 qa quotaAlarmer
29}
30
31type quotaAlarmer struct {
32 q etcdserver.Quota
33 a Alarmer
34 id types.ID
35}
36
37// check whether request satisfies the quota. If there is not enough space,
38// ignore request and raise the free space alarm.
39func (qa *quotaAlarmer) check(ctx context.Context, r interface{}) error {
40 if qa.q.Available(r) {
41 return nil
42 }
43 req := &pb.AlarmRequest{
44 MemberID: uint64(qa.id),
45 Action: pb.AlarmRequest_ACTIVATE,
46 Alarm: pb.AlarmType_NOSPACE,
47 }
48 qa.a.Alarm(ctx, req)
49 return rpctypes.ErrGRPCNoSpace
50}
51
52func NewQuotaKVServer(s *etcdserver.EtcdServer) pb.KVServer {
53 return &quotaKVServer{
54 NewKVServer(s),
55 quotaAlarmer{etcdserver.NewBackendQuota(s, "kv"), s, s.ID()},
56 }
57}
58
59func (s *quotaKVServer) Put(ctx context.Context, r *pb.PutRequest) (*pb.PutResponse, error) {
60 if err := s.qa.check(ctx, r); err != nil {
61 return nil, err
62 }
63 return s.KVServer.Put(ctx, r)
64}
65
66func (s *quotaKVServer) Txn(ctx context.Context, r *pb.TxnRequest) (*pb.TxnResponse, error) {
67 if err := s.qa.check(ctx, r); err != nil {
68 return nil, err
69 }
70 return s.KVServer.Txn(ctx, r)
71}
72
73type quotaLeaseServer struct {
74 pb.LeaseServer
75 qa quotaAlarmer
76}
77
78func (s *quotaLeaseServer) LeaseGrant(ctx context.Context, cr *pb.LeaseGrantRequest) (*pb.LeaseGrantResponse, error) {
79 if err := s.qa.check(ctx, cr); err != nil {
80 return nil, err
81 }
82 return s.LeaseServer.LeaseGrant(ctx, cr)
83}
84
85func NewQuotaLeaseServer(s *etcdserver.EtcdServer) pb.LeaseServer {
86 return &quotaLeaseServer{
87 NewLeaseServer(s),
88 quotaAlarmer{etcdserver.NewBackendQuota(s, "lease"), s, s.ID()},
89 }
90}