github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/id/id_gen.go (about)

     1  package id
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"time"
     7  
     8  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
     9  	"github.com/machinefi/w3bstream/pkg/depends/base/types/snowflake_id"
    10  	"github.com/machinefi/w3bstream/pkg/depends/x/contextx"
    11  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/must"
    12  )
    13  
    14  var (
    15  	start, _ = time.Parse(time.RFC3339, "2022-10-13T22:10:24+08:00")
    16  	sff      = snowflake_id.NewSnowflakeFactory(12, 10, 1, start.Local())
    17  )
    18  
    19  type Generator interface {
    20  	ID() (uint64, error)
    21  }
    22  
    23  func FromIP(ip net.IP) (Generator, error) {
    24  	return sff.NewSnowflake(snowflake_id.WorkerIDFromIP(ip))
    25  }
    26  
    27  func FromLocalIP() (Generator, error) {
    28  	return sff.NewSnowflake(snowflake_id.WorkerIDFromLocalIP())
    29  }
    30  
    31  type keyGenerator struct{}
    32  
    33  func WithGenerator(ctx context.Context, g Generator) context.Context {
    34  	return contextx.WithValue(ctx, keyGenerator{}, g)
    35  }
    36  
    37  func WithGeneratorContext(g Generator) contextx.WithContext {
    38  	return func(ctx context.Context) context.Context {
    39  		return contextx.WithValue(ctx, keyGenerator{}, g)
    40  	}
    41  }
    42  
    43  func GeneratorFromContext(ctx context.Context) (Generator, bool) {
    44  	g, ok := ctx.Value(keyGenerator{}).(Generator)
    45  	return g, ok
    46  }
    47  
    48  func MustGeneratorFromContext(ctx context.Context) Generator {
    49  	g, ok := ctx.Value(keyGenerator{}).(Generator)
    50  	must.BeTrue(ok)
    51  	return g
    52  }
    53  
    54  type SFIDGenerator interface {
    55  	MustGenSFID() types.SFID
    56  	MustGenSFIDs(n int) types.SFIDs
    57  }
    58  
    59  type keySFIDGenerator struct{}
    60  
    61  func WithSFIDGenerator(ctx context.Context, g SFIDGenerator) context.Context {
    62  	return contextx.WithValue(ctx, keySFIDGenerator{}, g)
    63  }
    64  
    65  func WithSFIDGeneratorContext(g SFIDGenerator) contextx.WithContext {
    66  	return func(ctx context.Context) context.Context {
    67  		return contextx.WithValue(ctx, keySFIDGenerator{}, g)
    68  	}
    69  }
    70  
    71  func SFIDGeneratorFromContext(ctx context.Context) (SFIDGenerator, bool) {
    72  	g, ok := ctx.Value(keySFIDGenerator{}).(SFIDGenerator)
    73  	return g, ok
    74  }
    75  
    76  func MustSFIDGeneratorFromContext(ctx context.Context) SFIDGenerator {
    77  	g, ok := ctx.Value(keySFIDGenerator{}).(SFIDGenerator)
    78  	must.BeTrue(ok)
    79  	return g
    80  }