1/*
2Copyright 2012 Google Inc.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8     http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17package groupcache
18
19import (
20	"errors"
21
22	"github.com/golang/protobuf/proto"
23)
24
25// A Sink receives data from a Get call.
26//
27// Implementation of Getter must call exactly one of the Set methods
28// on success.
29type Sink interface {
30	// SetString sets the value to s.
31	SetString(s string) error
32
33	// SetBytes sets the value to the contents of v.
34	// The caller retains ownership of v.
35	SetBytes(v []byte) error
36
37	// SetProto sets the value to the encoded version of m.
38	// The caller retains ownership of m.
39	SetProto(m proto.Message) error
40
41	// view returns a frozen view of the bytes for caching.
42	view() (ByteView, error)
43}
44
45func cloneBytes(b []byte) []byte {
46	c := make([]byte, len(b))
47	copy(c, b)
48	return c
49}
50
51func setSinkView(s Sink, v ByteView) error {
52	// A viewSetter is a Sink that can also receive its value from
53	// a ByteView. This is a fast path to minimize copies when the
54	// item was already cached locally in memory (where it's
55	// cached as a ByteView)
56	type viewSetter interface {
57		setView(v ByteView) error
58	}
59	if vs, ok := s.(viewSetter); ok {
60		return vs.setView(v)
61	}
62	if v.b != nil {
63		return s.SetBytes(v.b)
64	}
65	return s.SetString(v.s)
66}
67
68// StringSink returns a Sink that populates the provided string pointer.
69func StringSink(sp *string) Sink {
70	return &stringSink{sp: sp}
71}
72
73type stringSink struct {
74	sp *string
75	v  ByteView
76	// TODO(bradfitz): track whether any Sets were called.
77}
78
79func (s *stringSink) view() (ByteView, error) {
80	// TODO(bradfitz): return an error if no Set was called
81	return s.v, nil
82}
83
84func (s *stringSink) SetString(v string) error {
85	s.v.b = nil
86	s.v.s = v
87	*s.sp = v
88	return nil
89}
90
91func (s *stringSink) SetBytes(v []byte) error {
92	return s.SetString(string(v))
93}
94
95func (s *stringSink) SetProto(m proto.Message) error {
96	b, err := proto.Marshal(m)
97	if err != nil {
98		return err
99	}
100	s.v.b = b
101	*s.sp = string(b)
102	return nil
103}
104
105// ByteViewSink returns a Sink that populates a ByteView.
106func ByteViewSink(dst *ByteView) Sink {
107	if dst == nil {
108		panic("nil dst")
109	}
110	return &byteViewSink{dst: dst}
111}
112
113type byteViewSink struct {
114	dst *ByteView
115
116	// if this code ever ends up tracking that at least one set*
117	// method was called, don't make it an error to call set
118	// methods multiple times. Lorry's payload.go does that, and
119	// it makes sense. The comment at the top of this file about
120	// "exactly one of the Set methods" is overly strict. We
121	// really care about at least once (in a handler), but if
122	// multiple handlers fail (or multiple functions in a program
123	// using a Sink), it's okay to re-use the same one.
124}
125
126func (s *byteViewSink) setView(v ByteView) error {
127	*s.dst = v
128	return nil
129}
130
131func (s *byteViewSink) view() (ByteView, error) {
132	return *s.dst, nil
133}
134
135func (s *byteViewSink) SetProto(m proto.Message) error {
136	b, err := proto.Marshal(m)
137	if err != nil {
138		return err
139	}
140	*s.dst = ByteView{b: b}
141	return nil
142}
143
144func (s *byteViewSink) SetBytes(b []byte) error {
145	*s.dst = ByteView{b: cloneBytes(b)}
146	return nil
147}
148
149func (s *byteViewSink) SetString(v string) error {
150	*s.dst = ByteView{s: v}
151	return nil
152}
153
154// ProtoSink returns a sink that unmarshals binary proto values into m.
155func ProtoSink(m proto.Message) Sink {
156	return &protoSink{
157		dst: m,
158	}
159}
160
161type protoSink struct {
162	dst proto.Message // authoritative value
163	typ string
164
165	v ByteView // encoded
166}
167
168func (s *protoSink) view() (ByteView, error) {
169	return s.v, nil
170}
171
172func (s *protoSink) SetBytes(b []byte) error {
173	err := proto.Unmarshal(b, s.dst)
174	if err != nil {
175		return err
176	}
177	s.v.b = cloneBytes(b)
178	s.v.s = ""
179	return nil
180}
181
182func (s *protoSink) SetString(v string) error {
183	b := []byte(v)
184	err := proto.Unmarshal(b, s.dst)
185	if err != nil {
186		return err
187	}
188	s.v.b = b
189	s.v.s = ""
190	return nil
191}
192
193func (s *protoSink) SetProto(m proto.Message) error {
194	b, err := proto.Marshal(m)
195	if err != nil {
196		return err
197	}
198	// TODO(bradfitz): optimize for same-task case more and write
199	// right through? would need to document ownership rules at
200	// the same time. but then we could just assign *dst = *m
201	// here. This works for now:
202	err = proto.Unmarshal(b, s.dst)
203	if err != nil {
204		return err
205	}
206	s.v.b = b
207	s.v.s = ""
208	return nil
209}
210
211// AllocatingByteSliceSink returns a Sink that allocates
212// a byte slice to hold the received value and assigns
213// it to *dst. The memory is not retained by groupcache.
214func AllocatingByteSliceSink(dst *[]byte) Sink {
215	return &allocBytesSink{dst: dst}
216}
217
218type allocBytesSink struct {
219	dst *[]byte
220	v   ByteView
221}
222
223func (s *allocBytesSink) view() (ByteView, error) {
224	return s.v, nil
225}
226
227func (s *allocBytesSink) setView(v ByteView) error {
228	if v.b != nil {
229		*s.dst = cloneBytes(v.b)
230	} else {
231		*s.dst = []byte(v.s)
232	}
233	s.v = v
234	return nil
235}
236
237func (s *allocBytesSink) SetProto(m proto.Message) error {
238	b, err := proto.Marshal(m)
239	if err != nil {
240		return err
241	}
242	return s.setBytesOwned(b)
243}
244
245func (s *allocBytesSink) SetBytes(b []byte) error {
246	return s.setBytesOwned(cloneBytes(b))
247}
248
249func (s *allocBytesSink) setBytesOwned(b []byte) error {
250	if s.dst == nil {
251		return errors.New("nil AllocatingByteSliceSink *[]byte dst")
252	}
253	*s.dst = cloneBytes(b) // another copy, protecting the read-only s.v.b view
254	s.v.b = b
255	s.v.s = ""
256	return nil
257}
258
259func (s *allocBytesSink) SetString(v string) error {
260	if s.dst == nil {
261		return errors.New("nil AllocatingByteSliceSink *[]byte dst")
262	}
263	*s.dst = []byte(v)
264	s.v.b = nil
265	s.v.s = v
266	return nil
267}
268
269// TruncatingByteSliceSink returns a Sink that writes up to len(*dst)
270// bytes to *dst. If more bytes are available, they're silently
271// truncated. If fewer bytes are available than len(*dst), *dst
272// is shrunk to fit the number of bytes available.
273func TruncatingByteSliceSink(dst *[]byte) Sink {
274	return &truncBytesSink{dst: dst}
275}
276
277type truncBytesSink struct {
278	dst *[]byte
279	v   ByteView
280}
281
282func (s *truncBytesSink) view() (ByteView, error) {
283	return s.v, nil
284}
285
286func (s *truncBytesSink) SetProto(m proto.Message) error {
287	b, err := proto.Marshal(m)
288	if err != nil {
289		return err
290	}
291	return s.setBytesOwned(b)
292}
293
294func (s *truncBytesSink) SetBytes(b []byte) error {
295	return s.setBytesOwned(cloneBytes(b))
296}
297
298func (s *truncBytesSink) setBytesOwned(b []byte) error {
299	if s.dst == nil {
300		return errors.New("nil TruncatingByteSliceSink *[]byte dst")
301	}
302	n := copy(*s.dst, b)
303	if n < len(*s.dst) {
304		*s.dst = (*s.dst)[:n]
305	}
306	s.v.b = b
307	s.v.s = ""
308	return nil
309}
310
311func (s *truncBytesSink) SetString(v string) error {
312	if s.dst == nil {
313		return errors.New("nil TruncatingByteSliceSink *[]byte dst")
314	}
315	n := copy(*s.dst, v)
316	if n < len(*s.dst) {
317		*s.dst = (*s.dst)[:n]
318	}
319	s.v.b = nil
320	s.v.s = v
321	return nil
322}
323