Skip to content

Commit

Permalink
Refactor main.go to remove unused dependencies and task server initia…
Browse files Browse the repository at this point in the history 
…lization
  • Loading branch information
@cryptofyre
cryptofyre committed Sep 17, 2024
1 parent 09cfa27 commit 8aedd9b
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 140 deletions.
161 changes: 85 additions & 76 deletions main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ package main
import (
"crypto/md5"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"net/url"
Expand All @@ -14,16 +13,12 @@ import (
"time"

"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/hibiken/asynq"
"github.com/sirupsen/logrus"
ffmpeg "github.com/u2takey/ffmpeg-go"
)

var (
logger *logrus.Logger
redisClient *redis.Client
taskClient *asynq.Client
cacheDir string
artistSquares string
icloudArt string
Expand Down Expand Up @@ -59,14 +54,6 @@ func init() {
ffmpeg.LogCompiledCommand = false

ensureDirectories()

// Initialize Redis client
redisClient = redis.NewClient(&redis.Options{
Addr: "10.10.79.15:6379",
})

// Initialize Asynq client
taskClient = asynq.NewClient(asynq.RedisClientOpt{Addr: "10.10.79.15:6379"})
}

func ensureDirectories() {
Expand Down Expand Up @@ -95,31 +82,6 @@ func main() {
r.POST("/artwork/icloud", generateICloudArt)
r.GET("/artwork/icloud/:key", getICloudArt)

// Initialize task server
srv := asynq.NewServer(
asynq.RedisClientOpt{Addr: "10.10.79.15:6379"},
asynq.Config{
Concurrency: 10,
Queues: map[string]int{
"critical": 6,
"default": 3,
"low": 1,
},
},
)

// Register task handlers
mux := asynq.NewServeMux()
mux.HandleFunc(TypeCreateArtistSquare, HandleCreateArtistSquareTask)
mux.HandleFunc(TypeCreateICloudArt, HandleCreateICloudArtTask)

// Start task server
go func() {
if err := srv.Run(mux); err != nil {
logger.Fatalf("Could not run task server: %v", err)
}
}()

// Start server
if err := r.Run(":3000"); err != nil {
logger.Fatal("Failed to start server: ", err)
Expand Down Expand Up @@ -281,32 +243,57 @@ func generateArtistSquare(c *gin.Context) {
return
}

payload, err := json.Marshal(CreateArtistSquarePayload{
ImageURLs: request.ImageURLs,
Key: key,
JobID: generateKey(fmt.Sprintf("%s%s", key, request.ImageURLs[0])),
})
// Create a channel to receive the result
resultChan := make(chan error)

// Start a goroutine to generate the artist square
go func() {
err := generateArtistSquareAsync(request.ImageURLs, key, squarePath)
resultChan <- err
}()

// Wait for the goroutine to complete or timeout
select {
case err := <-resultChan:
if err != nil {
logger.Errorf("Failed to generate artist square: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate artist square"})
} else {
c.JSON(http.StatusOK, gin.H{
"key": key,
"message": "Artist square has been generated",
"url": fmt.Sprintf("https://art.cider.sh/artwork/artist-square/%s.jpg", key),
})
}
case <-time.After(30 * time.Second): // Adjust timeout as needed
c.JSON(http.StatusAccepted, gin.H{
"key": key,
"message": "Artist square is still being processed. Please check back later.",
"url": fmt.Sprintf("https://art.cider.sh/artwork/artist-square/%s.jpg", key),
})
}
}

func generateArtistSquareAsync(imageURLs []string, key, squarePath string) error {
images, err := downloadImages(imageURLs)
if err != nil {
logger.Errorf("Could not marshal task payload: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create task"})
return
logger.Errorf("Failed to download images: %v", err)
return fmt.Errorf("failed to download images: %w", err)
}

task := asynq.NewTask(TypeCreateArtistSquare, payload)
info, err := taskClient.Enqueue(task)
square, err := createArtistSquare(images)
if err != nil {
logger.Errorf("Could not enqueue task: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not enqueue task"})
return
logger.Errorf("Failed to create artist square: %v", err)
return fmt.Errorf("failed to create artist square: %w", err)
}

logger.Infof("Enqueued task: id=%s queue=%s", info.ID, info.Queue)
if err := saveJPEG(square, squarePath); err != nil {
logger.Errorf("Failed to save artist square: %v", err)
return fmt.Errorf("failed to save artist square: %w", err)
}

c.JSON(http.StatusAccepted, gin.H{
"key": key,
"message": "Artist square is being processed",
"url": fmt.Sprintf("https://art.cider.sh/artwork/artist-square/%s.jpg", key),
})
logger.Infof("Artist square created and saved successfully for key %s", key)
return nil
}

func getArtistSquare(c *gin.Context) {
Expand Down Expand Up @@ -353,32 +340,54 @@ func generateICloudArt(c *gin.Context) {
return
}

payload, err := json.Marshal(CreateICloudArtPayload{
ImageURL: request.ImageURL,
Key: key,
JobID: generateKey(fmt.Sprintf("%s%s", key, request.ImageURL)),
})
// Create a channel to receive the result
resultChan := make(chan error)

// Start a goroutine to generate the iCloud art
go func() {
err := generateICloudArtAsync(request.ImageURL, key, iCloudPath)
resultChan <- err
}()

// Wait for the goroutine to complete or timeout
select {
case err := <-resultChan:
if err != nil {
logger.Errorf("Failed to generate iCloud art: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate iCloud art"})
} else {
c.JSON(http.StatusOK, gin.H{
"key": key,
"message": "iCloud art has been generated",
"url": fmt.Sprintf("https://art.cider.sh/artwork/icloud/%s.jpg", key),
})
}
case <-time.After(30 * time.Second): // Adjust timeout as needed
c.JSON(http.StatusAccepted, gin.H{
"key": key,
"message": "iCloud art is still being processed. Please check back later.",
"url": fmt.Sprintf("https://art.cider.sh/artwork/icloud/%s.jpg", key),
})
}
}

func generateICloudArtAsync(imageURL, key, iCloudPath string) error {
img, err := downloadImage(imageURL)
if err != nil {
logger.Errorf("Could not marshal task payload: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create task"})
return
return fmt.Errorf("failed to download image: %w", err)
}

task := asynq.NewTask(TypeCreateICloudArt, payload)
info, err := taskClient.Enqueue(task)
iCloudImg, err := createICloudArt(img)
if err != nil {
logger.Errorf("Could not enqueue task: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not enqueue task"})
return
return fmt.Errorf("failed to create iCloud art: %w", err)
}

logger.Infof("Enqueued task: id=%s queue=%s", info.ID, info.Queue)
if err := saveJPEG(iCloudImg, iCloudPath); err != nil {
return fmt.Errorf("failed to save iCloud art: %w", err)
}

c.JSON(http.StatusAccepted, gin.H{
"key": key,
"message": "iCloud art is being processed",
"url": fmt.Sprintf("https://art.cider.sh/artwork/icloud/%s.jpg", key),
})
logger.Infof("iCloud art created and saved successfully for key %s", key)
return nil
}

func getICloudArt(c *gin.Context) {
Expand Down
64 changes: 0 additions & 64 deletions task.go
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
package main

import (
"context"
"encoding/json"
"fmt"
"image"
"image/draw"
Expand All @@ -11,7 +9,6 @@ import (
"os"
"path/filepath"

"github.com/hibiken/asynq"
"github.com/nfnt/resize"
"golang.org/x/image/webp"
)
Expand Down Expand Up @@ -40,38 +37,6 @@ type CreateICloudArtPayload struct {
JobID string `json:"job_id"`
}

func HandleCreateArtistSquareTask(ctx context.Context, t *asynq.Task) error {
var p CreateArtistSquarePayload
if err := json.Unmarshal(t.Payload(), &p); err != nil {
return fmt.Errorf("json.Unmarshal failed: %v: %w", err, asynq.SkipRetry)
}

squarePath := filepath.Join(artistSquares, fmt.Sprintf("%s.jpg", p.Key))

if _, err := os.Stat(squarePath); err == nil {
return nil
}

images, err := downloadImages(p.ImageURLs)
if err != nil {
logger.Errorf("Job %s: Failed to download images: %v", p.JobID, err)
return fmt.Errorf("failed to download images: %w", err)
}

square, err := createArtistSquare(images)
if err != nil {
logger.Errorf("Job %s: Failed to create artist square: %v", p.JobID, err)
return fmt.Errorf("failed to create artist square: %w", err)
}
if err := saveJPEG(square, squarePath); err != nil {
logger.Errorf("Job %s: Failed to save artist square: %v", p.JobID, err)
return fmt.Errorf("failed to save artist square: %w", err)
}

logger.Infof("Job %s: Artist square created and saved successfully for key %s", p.JobID, p.Key)
return nil
}

func downloadImages(urls []string) ([]image.Image, error) {
images := make([]image.Image, len(urls))
for i, url := range urls {
Expand Down Expand Up @@ -154,35 +119,6 @@ func saveJPEG(img image.Image, path string) error {
return jpeg.Encode(f, img, &jpeg.Options{Quality: 90})
}

func HandleCreateICloudArtTask(ctx context.Context, t *asynq.Task) error {
var p CreateICloudArtPayload
if err := json.Unmarshal(t.Payload(), &p); err != nil {
return fmt.Errorf("json.Unmarshal failed: %v: %w", err, asynq.SkipRetry)
}

iCloudPath := filepath.Join(icloudArt, fmt.Sprintf("%s.jpg", p.Key))

if _, err := os.Stat(iCloudPath); err == nil {
return nil
}

img, err := downloadImage(p.ImageURL)
if err != nil {
return fmt.Errorf("failed to download image: %w", err)
}

iCloudImg, err := createICloudArt(img)
if err != nil {
return fmt.Errorf("failed to create iCloud art: %w", err)
}

if err := saveJPEG(iCloudImg, iCloudPath); err != nil {
return fmt.Errorf("failed to save iCloud art: %w", err)
}

return nil
}

func downloadImage(url string) (image.Image, error) {
resp, err := http.Get(url)
if err != nil {
Expand Down

0 comments on commit 8aedd9b

Please sign in to comment.