Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move to xorshiro for rand generation #190

Merged
merged 1 commit into from
Nov 21, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/restate_context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -191,13 +191,13 @@ export interface Rand {
* Equivalent of JS `Math.random()` but deterministic; seeded by the invocation ID of the current invocation,
* each call will return a new pseudorandom float within the range [0,1)
*/
random(): number
random(): number;

/**
* Using the same random source and seed as random(), produce a UUID version 4 string. This is inherently predictable
* based on the invocation ID and should not be used in cryptographic contexts
*/
uuidv4(): string
uuidv4(): string;
}

// ----------------------------------------------------------------------------
Expand Down
2 changes: 1 addition & 1 deletion src/restate_context_impl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ import { rlog } from "./utils/logger";
import { Client, SendClient } from "./types/router";
import { RpcRequest, RpcResponse } from "./generated/proto/dynrpc";
import { requestFromArgs } from "./utils/assumpsions";
import {RandImpl} from "./utils/rand";
import { RandImpl } from "./utils/rand";

export enum CallContexType {
None,
Expand Down
7 changes: 5 additions & 2 deletions src/server/restate_lambda_handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ import { Message } from "../types/types";
import { StateMachine } from "../state_machine";
import { ensureError } from "../types/errors";
import { KeyedRouter, UnKeyedRouter } from "../public_api";
import {OUTPUT_STREAM_ENTRY_MESSAGE_TYPE} from "../types/protocol";
import { OUTPUT_STREAM_ENTRY_MESSAGE_TYPE } from "../types/protocol";

/**
* Creates an Restate entrypoint for services deployed on AWS Lambda and invoked
Expand Down Expand Up @@ -224,7 +224,10 @@ export class LambdaRestateServer extends BaseRestateServer {
let decodedEntries: Message[] | null = decodeLambdaBody(event.body);
const journalBuilder = new InvocationBuilder(method);
decodedEntries.forEach((e: Message) => journalBuilder.handleMessage(e));
const alreadyCompleted = decodedEntries.find((e: Message) => e.messageType === OUTPUT_STREAM_ENTRY_MESSAGE_TYPE) !== undefined
const alreadyCompleted =
decodedEntries.find(
(e: Message) => e.messageType === OUTPUT_STREAM_ENTRY_MESSAGE_TYPE
) !== undefined;
decodedEntries = null;

// set up and invoke the state machine
Expand Down
93 changes: 58 additions & 35 deletions src/utils/rand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,72 +12,95 @@
//! Some parts copied from https://github.com/uuidjs/uuid/blob/main/src/stringify.js
//! License MIT

import {Rand} from "../restate_context";
import {ErrorCodes, TerminalError} from "../types/errors";
import {CallContexType, RestateGrpcContextImpl} from "../restate_context_impl";
import {createHash} from "crypto";
import { Rand } from "../restate_context";
import { ErrorCodes, TerminalError } from "../types/errors";
import {
CallContexType,
RestateGrpcContextImpl,
} from "../restate_context_impl";
import { createHash } from "crypto";

export class RandImpl implements Rand {
private randstate64: bigint;
private randstate256: [bigint, bigint, bigint, bigint];

constructor(id: Buffer | bigint) {
if (typeof id == "bigint") {
this.randstate64 = id
} else {
constructor(id: Buffer | [bigint, bigint, bigint, bigint]) {
if (id instanceof Buffer) {
// hash the invocation ID, which is known to contain 74 bits of entropy
const hash = createHash('sha256')
.update(id)
.digest();

// seed using first 64 bits of the hash
this.randstate64 = hash.readBigUInt64LE(0);
const hash = createHash("sha256").update(id).digest();

this.randstate256 = [
hash.readBigUInt64LE(0),
hash.readBigUInt64LE(8),
hash.readBigUInt64LE(16),
hash.readBigUInt64LE(24),
];
} else {
this.randstate256 = id;
}
}

static U64_MASK = ((1n << 64n) - 1n)
static U64_MASK = (1n << 64n) - 1n;

// splitmix64
// https://prng.di.unimi.it/splitmix64.c - public domain
// xoshiro256++
// https://prng.di.unimi.it/xoshiro256plusplus.c - public domain
u64(): bigint {
this.randstate64 = (this.randstate64 + 0x9e3779b97f4a7c15n) & RandImpl.U64_MASK;
let next: bigint = this.randstate64;
next = ((next ^ (next >> 30n)) * 0xbf58476d1ce4e5b9n) & RandImpl.U64_MASK;
next = ((next ^ (next >> 27n)) * 0x94d049bb133111ebn) & RandImpl.U64_MASK;
next = next ^ (next >> 31n);
return next
const result: bigint =
(RandImpl.rotl(
(this.randstate256[0] + this.randstate256[3]) & RandImpl.U64_MASK,
23n
) +
this.randstate256[0]) &
RandImpl.U64_MASK;

const t: bigint = (this.randstate256[1] << 17n) & RandImpl.U64_MASK;

this.randstate256[2] ^= this.randstate256[0];
this.randstate256[3] ^= this.randstate256[1];
this.randstate256[1] ^= this.randstate256[2];
this.randstate256[0] ^= this.randstate256[3];

this.randstate256[2] ^= t;

this.randstate256[3] = RandImpl.rotl(this.randstate256[3], 45n);

return result;
}

static U53_MASK = ((1n << 53n) - 1n)
static rotl(x: bigint, k: bigint): bigint {
return ((x << k) & RandImpl.U64_MASK) | (x >> (64n - k));
}

checkContext() {
const context = RestateGrpcContextImpl.callContext.getStore();
if (context && context.type === CallContexType.SideEffect) {
throw new TerminalError(
`You may not call methods on Rand from within a side effect.`,
{errorCode: ErrorCodes.INTERNAL}
{ errorCode: ErrorCodes.INTERNAL }
);
}
}

static U53_MASK = (1n << 53n) - 1n;

public random(): number {
this.checkContext()
this.checkContext();

// first generate a uint in range [0,2^53), which can be mapped 1:1 to a float64 in [0,1)
const u53 = this.u64() & RandImpl.U53_MASK
const u53 = this.u64() & RandImpl.U53_MASK;
// then divide by 2^53, which will simply update the exponent
return Number(u53) / 2 ** 53
return Number(u53) / 2 ** 53;
}

public uuidv4(): string {
this.checkContext()
this.checkContext();

const buf = Buffer.alloc(16);
buf.writeBigUInt64LE(this.u64(), 0);
buf.writeBigUInt64LE(this.u64(), 8);
// Per 4.4, set bits for version and `clock_seq_hi_and_reserved`
buf[6] = (buf[6] & 0x0f) | 0x40;
buf[8] = (buf[8] & 0x3f) | 0x80;
return uuidStringify(buf)
return uuidStringify(buf);
}
}

Expand All @@ -102,16 +125,16 @@ function uuidStringify(arr: Buffer, offset = 0) {
byteToHex[arr[offset + 1]] +
byteToHex[arr[offset + 2]] +
byteToHex[arr[offset + 3]] +
'-' +
"-" +
byteToHex[arr[offset + 4]] +
byteToHex[arr[offset + 5]] +
'-' +
"-" +
byteToHex[arr[offset + 6]] +
byteToHex[arr[offset + 7]] +
'-' +
"-" +
byteToHex[arr[offset + 8]] +
byteToHex[arr[offset + 9]] +
'-' +
"-" +
byteToHex[arr[offset + 10]] +
byteToHex[arr[offset + 11]] +
byteToHex[arr[offset + 12]] +
Expand Down
12 changes: 8 additions & 4 deletions test/protoutils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,10 @@ export function startMessage(
return new Message(
START_MESSAGE_TYPE,
StartMessage.create({
id: Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"),
id: Buffer.from(
"f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2",
"hex"
),
debugId: "8xHx_cuYY_AAYvTQA7NfWm1RyBOd2IYsg",
knownEntries: knownEntries, // only used for the Lambda case. For bidi streaming, this will be imputed by the testdriver
stateMap: toStateEntries(state || []),
Expand Down Expand Up @@ -441,9 +444,10 @@ export function getAwakeableId(entryIndex: number): string {
const encodedEntryIndex = Buffer.alloc(4 /* Size of u32 */);
encodedEntryIndex.writeUInt32BE(entryIndex);

return Buffer.concat([Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"), encodedEntryIndex]).toString(
"base64url"
);
return Buffer.concat([
Buffer.from("f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2", "hex"),
encodedEntryIndex,
]).toString("base64url");
}

export function keyVal(key: string, value: any): Buffer[] {
Expand Down
Loading