Go学习 Day21 分布式原语之分布式队列

分布式队列和单机的队列是一样的，是一种先进先出的数据结构，利用etcd的一些方法可以实现分布式的队列，其主要方法包括：

1）NewQueue，func NewQueue(client *v3.Client, keyPrefix string) *Queue
2）入队，func (q *Queue) Enqueue(val string) error
3）出队，func (q *Queue) Enqueue(val string) error

demo：

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/coreos/etcd/clientv3"
)

const (
	queuePrefix = "/queue/"
	lockPrefix  = "/lock/"
)

type Queue struct {
	cli       *clientv3.Client
	queueName string
	lockName  string
}

func NewQueue(cli *clientv3.Client, name string) *Queue {
	q := &Queue{
		cli:       cli,
		queueName: queuePrefix + name,
		lockName:  lockPrefix + name,
	}
	return q
}

func (q *Queue) Push(data string) error {
	resp, err := q.cli.Grant(context.Background(), 10)
	if err != nil {
		return err
	}
	_, err = q.cli.Txn(context.Background()).
		If(clientv3.Compare(clientv3.Version(q.lockName), "=", 0)).
		Then(clientv3.OpPut(q.lockName, "locked", clientv3.WithLease(resp.ID))).
		Commit()
	if err != nil {
		return err
	}
	_, err = q.cli.Txn(context.Background()).
		If(clientv3.Compare(clientv3.Version(q.queueName), "=", 0)).
		Then(clientv3.OpPut(q.queueName, data)).
		Commit()
	if err != nil {
		q.cli.Delete(context.Background(), q.lockName)
		return err
	}
	q.cli.Delete(context.Background(), q.lockName)
	return nil
}

func (q *Queue) Pop() (string, error) {
	resp, err := q.cli.Grant(context.Background(), 10)
	if err != nil {
		return "", err
	}
	_, err = q.cli.Txn(context.Background()).
		If(clientv3.Compare(clientv3.Version(q.lockName), "=", 0)).
		Then(clientv3.OpPut(q.lockName, "locked", clientv3.WithLease(resp.ID))).
		Commit()
	if err != nil {
		return "", err
	}
	getResp, err := q.cli.Get(context.Background(), q.queueName)
	if err != nil {
		q.cli.Delete(context.Background(), q.lockName)
		return "", err
	}
	if len(getResp.Kvs) == 0 {
		q.cli.Delete(context.Background(), q.lockName)
		return "", nil
	}
	value := string(getResp.Kvs[0].Value)
	_, err = q.cli.Txn(context.Background()).
		If(clientv3.Compare(clientv3.Value(q.queueName), "=", value)).
		Then(clientv3.OpDelete(q.queueName)).
		Commit()
	if err != nil {
		q.cli.Delete(context.Background(), q.lockName)
		return "", err
	}
	q.cli.Delete(context.Background(), q.lockName)
	return value, nil
}

func main() {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   []string{"http://127.0.0.1:2379"},
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		log.Fatal(err)
	}
	defer cli.Close()

	q := NewQueue(cli, "example")

	err = q.Push("data for job 1")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Pushed data for job 1 to queue.")

	data, err := q.Pop()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Popped data from queue: %v\n", data)
}

此文章为2月Day21学习笔记，内容来源于极客时间《Go 并发编程实战课》