Skip to content

[BEAM-13628] Update SideInputCache to use full Transform and SideInputIDs as token information #16483

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.

Sign up for GitHub

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

Merged
lostluck merged 3 commits into from
Jan 11, 2022
Merged

[BEAM-13628] Update SideInputCache to use full Transform and SideInputIDs as token information #16483

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
eddabb2
refactor cache to use transform/side input IDs in keying
jrmccluskey Jan 11, 2022
d8fb9d5
Manage valid tokens as transform-sideinput-token structs
jrmccluskey Jan 11, 2022
c3d0097
Test case clean-up
jrmccluskey Jan 11, 2022
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
47 changes: 30 additions & 17 deletions sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,14 @@ import (

type token string

type cacheToken struct {
transformID string
sideInputID string
tok token
}

type cacheKey struct {
tok token
tok cacheToken
win string
key string
}
Expand All @@ -54,7 +60,7 @@ type SideInputCache struct {
mu sync.Mutex
cache map[cacheKey]exec.ReStream
idsToTokens map[string]token
validTokens map[token]int8 // Maps tokens to active bundle counts
validTokens map[cacheToken]int8 // Maps tokens to active bundle counts
metrics CacheMetrics
}

Expand All @@ -78,7 +84,7 @@ func (c *SideInputCache) Init(cap int) error {
}
c.cache = make(map[cacheKey]exec.ReStream, cap)
c.idsToTokens = make(map[string]token)
c.validTokens = make(map[token]int8)
c.validTokens = make(map[cacheToken]int8)
c.capacity = cap
c.metrics = CacheMetrics{}
c.enabled = true
Expand Down Expand Up @@ -113,11 +119,12 @@ func (c *SideInputCache) SetValidTokens(cacheTokens ...*fnpb.ProcessBundleReques
func (c *SideInputCache) setValidToken(transformID, sideInputID string, tok token) {
idKey := transformID + sideInputID
c.idsToTokens[idKey] = tok
count, ok := c.validTokens[tok]
fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok}
count, ok := c.validTokens[fullToken]
if !ok {
c.validTokens[tok] = 1
c.validTokens[fullToken] = 1
} else {
c.validTokens[tok] = count + 1
c.validTokens[fullToken] = count + 1
}
}

Expand All @@ -135,20 +142,24 @@ func (c *SideInputCache) CompleteBundle(cacheTokens ...*fnpb.ProcessBundleReques
if tok.GetUserState() != nil {
continue
}
s := tok.GetSideInput()
transformID := s.GetTransformId()
sideInputID := s.GetSideInputId()
t := token(tok.GetToken())
c.decrementTokenCount(t)
c.decrementTokenCount(transformID, sideInputID, t)
}
}

// decrementTokenCount decrements the validTokens entry for
// a given token by 1. Should only be called when completing
// a bundle.
func (c *SideInputCache) decrementTokenCount(tok token) {
count := c.validTokens[tok]
func (c *SideInputCache) decrementTokenCount(transformID, sideInputID string, tok token) {
fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok}
count := c.validTokens[fullToken]
if count == 1 {
delete(c.validTokens, tok)
delete(c.validTokens, fullToken)
} else {
c.validTokens[tok] = count - 1
c.validTokens[fullToken] = count - 1
}
}

Expand All @@ -159,11 +170,13 @@ func (c *SideInputCache) makeAndValidateToken(transformID, sideInputID string) (
if !ok {
return "", false
}
return tok, c.isValid(tok)
fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok}
return tok, c.isValid(fullToken)
}

func (c *SideInputCache) makeCacheKey(tok token, w, key []byte) cacheKey {
return cacheKey{tok: tok, win: string(w), key: string(key)}
func (c *SideInputCache) makeCacheKey(transformID, sideInputID string, tok token, w, key []byte) cacheKey {
fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok}
return cacheKey{tok: fullToken, win: string(w), key: string(key)}
}

// QueryCache takes a transform ID and side input ID and checking if a corresponding side
Expand All @@ -180,7 +193,7 @@ func (c *SideInputCache) QueryCache(ctx context.Context, transformID, sideInputI
if !ok {
return nil
}
ck := c.makeCacheKey(tok, win, key)
ck := c.makeCacheKey(transformID, sideInputID, tok, win, key)
// Check to see if cached
input, ok := c.cache[ck]
if !ok {
Expand Down Expand Up @@ -224,12 +237,12 @@ func (c *SideInputCache) SetCache(ctx context.Context, transformID, sideInputID
c.metrics.ReStreamErrors++
return input
}
ck := c.makeCacheKey(tok, win, key)
ck := c.makeCacheKey(transformID, sideInputID, tok, win, key)
c.cache[ck] = mat
return mat
}

func (c *SideInputCache) isValid(tok token) bool {
func (c *SideInputCache) isValid(tok cacheToken) bool {
count, ok := c.validTokens[tok]
// If the token is not known or not in use, return false
return ok && count > 0
Expand Down
54 changes: 30 additions & 24 deletions sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ func TestQueryCache_EmptyCase(t *testing.T) {
}
output := s.QueryCache(ctx, "side1", "transform1", win, key)
if output != nil {
t.Errorf("Cache hit when it should have missed, got %v", output)
t.Errorf("cache hit when it should have missed, got %v", output)
}
}

Expand All @@ -89,7 +89,7 @@ func TestSetCache_UncacheableCase(t *testing.T) {
s.SetCache(ctx, "t1", "s1", win, key, input)
output := s.QueryCache(ctx, "t1", "s1", win, key)
if output != nil {
t.Errorf("Cache hit when should have missed, got %v", output)
t.Errorf("cache hit when should have missed, got %v", output)
}
}

Expand Down Expand Up @@ -117,7 +117,7 @@ func TestSetCache_CacheableCase(t *testing.T) {
t.Errorf("failed to convert value to integer, got %v", getValue(output))
}
if val != 10 {
t.Errorf("element mismatch, expected 10, got %v", val)
t.Errorf("element mismatch, got %v, want 10", val)
}
}

Expand All @@ -133,6 +133,10 @@ func makeRequest(transformID, sideInputID string, t token) *fnpb.ProcessBundleRe
}
}

func makeCacheToken(transformID, sideInputID string, tok token) cacheToken {
return cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok}
}

func TestSetValidTokens(t *testing.T) {
inputs := []struct {
transformID string
Expand Down Expand Up @@ -169,19 +173,20 @@ func TestSetValidTokens(t *testing.T) {
}

s.SetValidTokens(tokens...)
if len(s.idsToTokens) != len(inputs) {
t.Errorf("Missing tokens, expected %v, got %v", len(inputs), len(s.idsToTokens))
if got, want := len(s.idsToTokens), len(inputs); got != want {
t.Errorf("got %d tokens, want %d", got, want)
}

for i, input := range inputs {
fullTok := makeCacheToken(input.transformID, input.sideInputID, input.tok)
// Check that the token is in the valid list
if !s.isValid(input.tok) {
t.Errorf("error in input %v, token %v is not valid", i, input.tok)
if !s.isValid(fullTok) {
t.Errorf("error in input %v, token %v is not valid", i, fullTok)
}
// Check that the mapping of IDs to tokens is correct
mapped := s.idsToTokens[input.transformID+input.sideInputID]
if mapped != input.tok {
t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tok, mapped)
got := s.idsToTokens[input.transformID+input.sideInputID]
if got != input.tok {
t.Errorf("got token %v for element %d, want %v", got, i, input.tok)
}
}
}
Expand Down Expand Up @@ -221,21 +226,22 @@ func TestSetValidTokens_ClearingBetween(t *testing.T) {
s.SetValidTokens(tok)

// Check that the token is in the valid list
if !s.isValid(input.tk) {
fullTok := makeCacheToken(input.transformID, input.sideInputID, input.tk)
if !s.isValid(fullTok) {
t.Errorf("error in input %v, token %v is not valid", i, input.tk)
}
// Check that the mapping of IDs to tokens is correct
mapped := s.idsToTokens[input.transformID+input.sideInputID]
if mapped != input.tk {
t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tk, mapped)
got := s.idsToTokens[input.transformID+input.sideInputID]
if got != input.tk {
t.Errorf("got token %v for element %d, want %v", got, i, input.tk)
}

s.CompleteBundle(tok)
}

for k := range s.validTokens {
if s.validTokens[k] != 0 {
t.Errorf("token count mismatch for token %v, expected 0, got %v", k, s.validTokens[k])
if got, want := s.validTokens[k], int8(0); got != want {
t.Errorf("got %d total valid tokens, want %d", got, want)
}
}
}
Expand All @@ -262,11 +268,11 @@ func TestSetCache_Eviction(t *testing.T) {
s.SetValidTokens(tokTwo)
s.SetCache(ctx, "t2", "s2", win, key, inTwo)

if len(s.cache) != 1 {
t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache))
if got, want := len(s.cache), 1; got != want {
t.Errorf("got %d elements in cache, want %d", got, want)
}
if s.metrics.Evictions != 1 {
t.Errorf("number evictions incorrect, expected 1, got %v", s.metrics.Evictions)
if got, want := s.metrics.Evictions, int64(1); got != want {
t.Errorf("got %d evictions, want %d", got, want)
}
}

Expand All @@ -291,10 +297,10 @@ func TestSetCache_EvictionFailure(t *testing.T) {
// Should fail to evict because the first token is still valid
s.SetCache(ctx, "t2", "s2", win, key, inTwo)
// Cache should not exceed size 1
if len(s.cache) != 1 {
t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache))
if got, want := len(s.cache), 1; got != want {
t.Errorf("got cache size of %d, want %d", got, want)
}
if s.metrics.InUseEvictions != 1 {
t.Errorf("number of failed evicition calls incorrect, expected 1, got %v", s.metrics.InUseEvictions)
if got, want := s.metrics.InUseEvictions, int64(1); got != want {
t.Errorf("got %d in-use eviction calls, want %d", got, want)
}
}