Skip to content

Commit

Permalink
add payload decompression with zstd algo
Browse files Browse the repository at this point in the history
  • Loading branch information
MindHunter86 committed Oct 19, 2024
1 parent 375ac51 commit bdc9467
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 2 deletions.
5 changes: 5 additions & 0 deletions cmd/alice/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,11 @@ func flagsInitialization(expertMode bool) []cli.Flag {
Category: "Release randomizer",
Value: 0,
},
&cli.BoolFlag{
Name: "randomizer-redis-zstd-enable",
Category: "Release randomizer",
Usage: "enable redis payload decompression with zstd algo",
},
&cli.StringFlag{
Name: "randomizer-releaseskey",
Category: "Release randomizer",
Expand Down
35 changes: 33 additions & 2 deletions internal/anilibria/randomizer.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (

"github.com/anilibria/alice/internal/utils"
futils "github.com/gofiber/fiber/v2/utils"
"github.com/klauspost/compress/zstd"
"github.com/redis/go-redis/v9"
"github.com/rs/zerolog"
"github.com/urfave/cli/v2"
Expand All @@ -30,13 +31,20 @@ type Randomizer struct {
relUpdFreqErr time.Duration
relUpdFreqBoot time.Duration

decoder *zstd.Decoder

mu sync.RWMutex
releases []string
}

func New(c context.Context) *Randomizer {
cli := c.Value(utils.CKCliCtx).(*cli.Context)

var dec *zstd.Decoder
if cli.Bool("randomizer-redis-zstd-enable") {
dec, _ = zstd.NewReader(nil, zstd.WithDecoderConcurrency(0))
}

r := &Randomizer{
done: c.Done,
log: c.Value(utils.CKLogger).(*zerolog.Logger),
Expand All @@ -60,6 +68,8 @@ func New(c context.Context) *Randomizer {
relUpdFreqErr: cli.Duration("randomizer-update-frequency-onerror"),
relUpdFreqBoot: cli.Duration("randomizer-update-frequency-bootstrap"),

decoder: dec,

releases: make([]string, 0),
releasesKey: cli.String("randomizer-releaseskey"),
}
Expand Down Expand Up @@ -125,7 +135,13 @@ func (m *Randomizer) peekReleaseKeyChunks() (_ int, e error) {
return
}

return strconv.Atoi(res)
var dres []byte
if dres, e = m.decompressPayload(futils.UnsafeBytes(res)); e != nil {
m.log.Warn().Msg("an error occurred while decompress response redis response - " + e.Error())
return
}

return strconv.Atoi(futils.UnsafeString(dres))
}

func (m *Randomizer) lookupReleases() (_ []string, e error) {
Expand Down Expand Up @@ -167,8 +183,15 @@ func (m *Randomizer) lookupReleases() (_ []string, e error) {
continue
}

var dres []byte
if dres, e = m.decompressPayload(futils.UnsafeBytes(res)); e != nil {
m.log.Warn().Msg("an error occurred while decompress redis response - " + e.Error())
errs = append(errs, e.Error())
continue
}

var releasesChunk Releases
if e = json.Unmarshal(futils.UnsafeBytes(res), &releasesChunk); e != nil {
if e = json.Unmarshal(dres, &releasesChunk); e != nil {
m.log.Warn().Msg("an error occurred while unmarshal release chunk - " + e.Error())
errs = append(errs, e.Error())
continue
Expand Down Expand Up @@ -229,3 +252,11 @@ func (m *Randomizer) randomRelease() (_ string) {
r := rand.Intn(len(m.releases)) // skipcq: GSC-G404 math/rand is enoght here
return m.releases[r]
}

func (m *Randomizer) decompressPayload(payload []byte) ([]byte, error) {
if m.decoder == nil {
return payload, nil
}

return m.decoder.DecodeAll(payload, nil)
}

0 comments on commit bdc9467

Please sign in to comment.