hatchet-dev · abelanger5 · Jul 26, 2024 · Jul 23, 2024 · Jul 23, 2024 · Jul 23, 2024
diff --git a/examples/assignment-affinity/main.go b/examples/assignment-affinity/main.go
@@ -0,0 +1,38 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/joho/godotenv"
+
+	"github.com/hatchet-dev/hatchet/pkg/cmdutils"
+)
+
+type userCreateEvent struct {
+	Username string            `json:"username"`
+	UserID   string            `json:"user_id"`
+	Data     map[string]string `json:"data"`
+}
+
+type stepOneOutput struct {
+	Message string `json:"message"`
+}
+
+func main() {
+	err := godotenv.Load()
+	if err != nil {
+		panic(err)
+	}
+
+	ch := cmdutils.InterruptChan()
+	cleanup, err := run()
+	if err != nil {
+		panic(err)
+	}
+
+	<-ch
+
+	if err := cleanup(); err != nil {
+		panic(fmt.Errorf("cleanup() error = %v", err))
+	}
+}
diff --git a/examples/assignment-affinity/run.go b/examples/assignment-affinity/run.go
@@ -0,0 +1,106 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/hatchet-dev/hatchet/pkg/client"
+	"github.com/hatchet-dev/hatchet/pkg/client/types"
+	"github.com/hatchet-dev/hatchet/pkg/worker"
+)
+
+func run() (func() error, error) {
+	c, err := client.New()
+	if err != nil {
+		return nil, fmt.Errorf("error creating client: %w", err)
+	}
+
+	w, err := worker.NewWorker(
+		worker.WithClient(
+			c,
+		),
+		worker.WithLabels(map[string]interface{}{
+			"model":  "fancy-ai-model-v2",
+			"memory": 512,
+		}),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("error creating worker: %w", err)
+	}
+
+	err = w.RegisterWorkflow(
+		&worker.WorkflowJob{
+			On:          worker.Events("user:create:affinity"),
+			Name:        "affinity",
+			Description: "affinity",
+			Steps: []*worker.WorkflowStep{
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+
+					model := ctx.Worker().GetLabels()["model"]
+
+					if model != "fancy-ai-model-v3" {
+						ctx.Worker().UpsertLabels(map[string]interface{}{
+							"model": nil,
+						})
+						// Do something to load the model
+						ctx.Worker().UpsertLabels(map[string]interface{}{
+							"model": "fancy-ai-model-v3",
+						})
+					}
+
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).
+					SetName("step-one").
+					SetDesiredLabels(map[string]*types.DesiredWorkerLabel{
+						"model": {
+							Value:  "fancy-ai-model-v3",
+							Weight: 10,
+						},
+						"memory": {
+							Value:      512,
+							Required:   true,
+							Comparator: types.ComparatorPtr(types.WorkerLabelComparator_GREATER_THAN),
+						},
+					}),
+			},
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("error registering workflow: %w", err)
+	}
+
+	go func() {
+		log.Printf("pushing event")
+
+		testEvent := userCreateEvent{
+			Username: "echo-test",
+			UserID:   "1234",
+			Data: map[string]string{
+				"test": "test",
+			},
+		}
+
+		// push an event
+		err := c.Event().Push(
+			context.Background(),
+			"user:create:affinity",
+			testEvent,
+		)
+		if err != nil {
+			panic(fmt.Errorf("error pushing event: %w", err))
+		}
+
+		time.Sleep(10 * time.Second)
+	}()
+
+	cleanup, err := w.Start()
+	if err != nil {
+		return nil, fmt.Errorf("error starting worker: %w", err)
+	}
+
+	return cleanup, nil
+}
diff --git a/examples/assignment-sticky/main.go b/examples/assignment-sticky/main.go
@@ -0,0 +1,38 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/joho/godotenv"
+
+	"github.com/hatchet-dev/hatchet/pkg/cmdutils"
+)
+
+type userCreateEvent struct {
+	Username string            `json:"username"`
+	UserID   string            `json:"user_id"`
+	Data     map[string]string `json:"data"`
+}
+
+type stepOneOutput struct {
+	Message string `json:"message"`
+}
+
+func main() {
+	err := godotenv.Load()
+	if err != nil {
+		panic(err)
+	}
+
+	ch := cmdutils.InterruptChan()
+	cleanup, err := run()
+	if err != nil {
+		panic(err)
+	}
+
+	<-ch
+
+	if err := cleanup(); err != nil {
+		panic(fmt.Errorf("cleanup() error = %v", err))
+	}
+}
diff --git a/examples/assignment-sticky/run.go b/examples/assignment-sticky/run.go
@@ -0,0 +1,99 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/hatchet-dev/hatchet/pkg/client"
+	"github.com/hatchet-dev/hatchet/pkg/client/types"
+	"github.com/hatchet-dev/hatchet/pkg/worker"
+)
+
+func run() (func() error, error) {
+	c, err := client.New()
+	if err != nil {
+		return nil, fmt.Errorf("error creating client: %w", err)
+	}
+
+	w, err := worker.NewWorker(
+		worker.WithClient(
+			c,
+		),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("error creating worker: %w", err)
+	}
+
+	err = w.RegisterWorkflow(
+		&worker.WorkflowJob{
+			On:             worker.Events("user:create:sticky"),
+			Name:           "sticky",
+			Description:    "sticky",
+			StickyStrategy: types.StickyStrategyPtr(types.StickyStrategy_SOFT),
+			Steps: []*worker.WorkflowStep{
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+
+					sticky := true
+
+					_, err = ctx.SpawnWorkflow("step-one", nil, &worker.SpawnWorkflowOpts{
+						Sticky: &sticky,
+					})
+
+					if err != nil {
+						return nil, fmt.Errorf("error spawning workflow: %w", err)
+					}
+
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-one"),
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-two"),
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-three").AddParents("step-one", "step-two"),
+			},
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("error registering workflow: %w", err)
+	}
+
+	go func() {
+		log.Printf("pushing event")
+
+		testEvent := userCreateEvent{
+			Username: "echo-test",
+			UserID:   "1234",
+			Data: map[string]string{
+				"test": "test",
+			},
+		}
+
+		// push an event
+		err := c.Event().Push(
+			context.Background(),
+			"user:create:sticky",
+			testEvent,
+		)
+		if err != nil {
+			panic(fmt.Errorf("error pushing event: %w", err))
+		}
+
+		time.Sleep(10 * time.Second)
+	}()
+
+	cleanup, err := w.Start()
+	if err != nil {
+		return nil, fmt.Errorf("error starting worker: %w", err)
+	}
+
+	return cleanup, nil
+}
diff --git a/frontend/docs/pages/home/features/worker-assignment/sticky-assignment.mdx b/frontend/docs/pages/home/features/worker-assignment/sticky-assignment.mdx
@@ -70,7 +70,31 @@ const myWorkflow: Workflow = {
   </Tabs.Tab>
   <Tabs.Tab>
 ```go
-// TODO add go example
+	err = w.RegisterWorkflow(
+		&worker.WorkflowJob{
+			On:             worker.Events("user:create:sticky"),
+			Name:           "sticky",
+			Description:    "sticky",
+			StickyStrategy: types.StickyStrategyPtr(types.StickyStrategy_SOFT),
+			Steps: []*worker.WorkflowStep{
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-one"),
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-two"),
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+					return &stepOneOutput{
+						Message: ctx.Worker().ID(),
+					}, nil
+				}).SetName("step-three").AddParents("step-one", "step-two"),
+			},
+		},
+	)
 ```
   </Tabs.Tab>
 </Tabs>
@@ -143,7 +167,33 @@ const parentWorkflow: Workflow = {
   </Tabs.Tab>
   <Tabs.Tab>
 ```go
-// TODO go example
+	err = w.RegisterWorkflow(
+		&worker.WorkflowJob{
+			On:             worker.Events("user:create:sticky"),
+			Name:           "sticky",
+			Description:    "sticky",
+			Steps: []*worker.WorkflowStep{
+				worker.Fn(func(ctx worker.HatchetContext) (result *stepOneOutput, err error) {
+
+    				sticky := true
+
+    				_, err = ctx.SpawnWorkflow("sticky-child", nil, &worker.SpawnWorkflowOpts{
+    					Sticky: &sticky,
+    				})
+
+    				if err != nil {
+    					return nil, fmt.Errorf("error spawning workflow: %w", err)
+    				}
+
+    				return &stepOneOutput{
+    					Message: ctx.Worker().ID(),
+    				}, nil
+    			}).SetName("step-one"),
+    		},
+    	},
+    )
+
 ```
   </Tabs.Tab>
 </Tabs>
+```