1// Copyright (C) 2021 The Android Open Source Project
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
15import {assertTrue} from './logging';
16
17// This class is the TypeScript equivalent of the identically-named C++ class in
18// //protozero/proto_ring_buffer.h. See comments in that header for a detailed
19// description. The architecture is identical.
20
21const kGrowBytes = 128 * 1024;
22const kMaxMsgSize = 64 * 1024 * 1024;
23
24export class ProtoRingBuffer {
25  private buf = new Uint8Array(kGrowBytes);
26  private fastpath?: Uint8Array;
27  private rd = 0;
28  private wr = 0;
29
30  // The caller must call ReadMessage() after each append() call.
31  // The |data| might be either copied in the internal ring buffer or returned
32  // (% subarray()) to the next ReadMessage() call.
33  append(data: Uint8Array) {
34    assertTrue(this.wr <= this.buf.length);
35    assertTrue(this.rd <= this.wr);
36
37    // If the last call to ReadMessage() consumed all the data in the buffer and
38    // there are no incomplete messages pending, restart from the beginning
39    // rather than keep ringing. This is the most common case.
40    if (this.rd === this.wr) {
41      this.rd = this.wr = 0;
42    }
43
44    // The caller is expected to issue a ReadMessage() after each append().
45    const dataLen = data.length;
46    if (dataLen === 0) return;
47    assertTrue(this.fastpath === undefined);
48    if (this.rd === this.wr) {
49      const msg = ProtoRingBuffer.tryReadMessage(data, 0, dataLen);
50      if (msg !== undefined &&
51          ((msg.byteOffset + msg.length) === (data.byteOffset + dataLen))) {
52        // Fastpath: in many cases, the underlying stream will effectively
53        // preserve the atomicity of messages for most small messages.
54        // In this case we can avoid the extra buffer roundtrip and return the
55        // original array (actually a subarray that skips the proto header).
56        // The next call to ReadMessage() will return this.
57        this.fastpath = msg;
58        return;
59      }
60    }
61
62    let avail = this.buf.length - this.wr;
63    if (dataLen > avail) {
64      // This whole section should be hit extremely rarely.
65
66      // Try first just recompacting the buffer by moving everything to the
67      // left. This can happen if we received "a message and a bit" on each
68      // append() call.
69      this.buf.copyWithin(0, this.rd, this.wr);
70      avail += this.rd;
71      this.wr -= this.rd;
72      this.rd = 0;
73      if (dataLen > avail) {
74        // Still not enough, expand the buffer.
75        let newSize = this.buf.length;
76        while (dataLen > newSize - this.wr) {
77          newSize += kGrowBytes;
78        }
79        assertTrue(newSize <= kMaxMsgSize * 2);
80        const newBuf = new Uint8Array(newSize);
81        newBuf.set(this.buf);
82        this.buf = newBuf;
83        // No need to touch rd / wr.
84      }
85    }
86
87    // Append the received data at the end of the ring buffer.
88    this.buf.set(data, this.wr);
89    this.wr += dataLen;
90  }
91
92  // Tries to extract a message from the ring buffer. If there is no message,
93  // or if the current message is still incomplete, returns undefined.
94  // The caller is expected to call this in a loop until it returns undefined.
95  // Note that a single write to Append() can yield more than one message
96  // (see ProtoRingBufferTest.CoalescingStream in the unittest).
97  readMessage(): Uint8Array|undefined {
98    if (this.fastpath !== undefined) {
99      assertTrue(this.rd === this.wr);
100      const msg = this.fastpath;
101      this.fastpath = undefined;
102      return msg;
103    }
104    assertTrue(this.rd <= this.wr);
105    if (this.rd >= this.wr) {
106      return undefined;  // Completely empty.
107    }
108    const msg = ProtoRingBuffer.tryReadMessage(this.buf, this.rd, this.wr);
109    if (msg === undefined) return undefined;
110    assertTrue(msg.buffer === this.buf.buffer);
111    assertTrue(this.buf.byteOffset === 0);
112    this.rd = msg.byteOffset + msg.length;
113
114    // Deliberately returning a copy of the data with slice(). In various cases
115    // (streaming query response) the caller will hold onto the returned buffer.
116    // If we get to this point, |msg| is a view of the circular buffer that we
117    // will overwrite on the next calls to append().
118    return msg.slice();
119  }
120
121  private static tryReadMessage(
122      data: Uint8Array, dataStart: number, dataEnd: number): Uint8Array
123      |undefined {
124    assertTrue(dataEnd <= data.length);
125    let pos = dataStart;
126    if (pos >= dataEnd) return undefined;
127    const tag = data[pos++];  // Assume one-byte tag.
128    if (tag >= 0x80 || (tag & 0x07) !== 2 /* len delimited */) {
129      throw new Error(
130          `RPC framing error, unexpected tag ${tag} @ offset ${pos - 1}`);
131    }
132
133    let len = 0;
134    for (let shift = 0; /* no check */; shift += 7) {
135      if (pos >= dataEnd) {
136        return undefined;  // Not enough data to read varint.
137      }
138      const val = data[pos++];
139      len |= ((val & 0x7f) << shift) >>> 0;
140      if (val < 0x80) break;
141    }
142
143    if (len >= kMaxMsgSize) {
144      throw new Error(
145          `RPC framing error, message too large (${len} > ${kMaxMsgSize}`);
146    }
147    const end = pos + len;
148    if (end > dataEnd) return undefined;
149
150    // This is a subarray() and not a slice() because in the |fastpath| case
151    // we want to just return the original buffer pushed by append().
152    // In the slow-path (ring-buffer) case, the readMessage() above will create
153    // a copy via slice() before returning it.
154    return data.subarray(pos, end);
155  }
156}
157