6.5840-mapreduce

package mr

import (
	"log"
	"net"
	"os"
	"net/rpc"
	"net/http"
	"time"
	"sync"
)

type Coordinator struct {
	// Your definitions here.

	// 一把大锁锁所有
	mtx sync.Mutex
	
	// MAP阶段涉及的状态
	toMapFiles map[int]string
	mappingFiles map[int]string
	mapDone map[int]bool
	mapPhaseDone bool
	mapDoneCnt int
	nMap int

	// REDUCE阶段涉及的状态
	toReduce map[int]bool
	reducing map[int]bool
	reduceDone map[int]bool
	reducePhaseDone bool
	reduceDoneCnt int
	nReduce int
}

// Your code here -- RPC handlers for the worker to call.
// 获取当前状态
func (c *Coordinator) Status(args *struct{}, reply *StatusReply) error {
	if !c.mapPhaseDone {
		reply.Status = 0
	} else if !c.reducePhaseDone {
		reply.Status = 1
	} else {
		reply.Status = 2
	}
	return nil
}

func (c *Coordinator) Map(args *struct{}, reply *MapReply) error {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	if c.mapPhaseDone || len(c.toMapFiles) == 0 {
		reply.MapperId = -1
		return nil
	}
	// 获取一个等待map的文件，将其从mappingFiles移除，加入到mappingFiles
	for mapperId, file := range c.toMapFiles {
		reply.MapperId = mapperId
		reply.NReduce = c.nReduce
		reply.Inputfile = file
		c.mappingFiles[mapperId] = file
		delete(c.toMapFiles, mapperId)
		break
	}

	// 等待10s之后，如果还没完成map任务，重新加入toMapFiles
	mid := reply.MapperId
	go func() {
        <-time.After(10 * time.Second)
		c.mtx.Lock()
		defer c.mtx.Unlock()
        if !c.mapDone[mid] {
			log.Printf("Map %d failed\n", mid) 
			c.toMapFiles[mid] = c.mappingFiles[mid]
			delete(c.mappingFiles, mid)
		}
    } ()

	return nil
}

func (c *Coordinator) MapDone(args *MapDoneArgs, reply *struct{}) error {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	c.mapDoneCnt++
	c.mapDone[args.MapperId] = true
	if c.mapDoneCnt == c.nMap {
		c.mapPhaseDone = true
	}
	delete(c.mappingFiles, args.MapperId)
	log.Printf("MapDone %d!", args.MapperId)
	
	return nil
}

func (c *Coordinator) Reduce(args *struct{}, reply *ReduceReply) error {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	if c.reducePhaseDone || len(c.toReduce) == 0 {
		reply.ReducerId = -1
		return nil
	}
	// 获取一个等待map的文件，将其从mappingFiles移除，加入到mappingFiles
	for reducerId, _ := range c.toReduce {
		reply.ReducerId = reducerId
		reply.NReduce = c.nReduce
		reply.NMap = c.nMap
		c.reducing[reducerId] = true
		delete(c.toReduce, reducerId)
		break
	}

	// 等待10s之后，如果还没完成map任务，重新加入toMapFiles
	rid := reply.ReducerId
	go func() {
        <-time.After(10 * time.Second)
		c.mtx.Lock()
		defer c.mtx.Unlock()
        if !c.reduceDone[rid] {
			log.Printf("Reduce %d failed\n", rid)
			c.toReduce[rid] = false
			delete(c.reducing, rid)
		}
    } ()

	return nil
}

func (c *Coordinator) ReduceDone(args *ReduceDoneArgs, reply *struct{}) error {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	c.reduceDoneCnt++
	c.reduceDone[args.ReducerId] = true
	if c.reduceDoneCnt == c.nReduce {
		c.reducePhaseDone = true
	}
	delete(c.reducing, args.ReducerId)
	log.Printf("ReduceDone %d!\n", args.ReducerId)
	return nil
}



//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {
	ret := c.mapPhaseDone && c.reducePhaseDone
	if ret {
		log.Println("Coordinator Done")
	}
	return ret
}


//
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
// 
// files是输入的文件
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{
		toMapFiles: make(map[int]string),
		mappingFiles: make(map[int]string),
		mapDone: make(map[int]bool),
		toReduce: make(map[int]bool),
		reducing: make(map[int]bool),
		reduceDone: make(map[int]bool),
		mapPhaseDone: false,
		reducePhaseDone: false,
		mapDoneCnt: 0,
		reduceDoneCnt: 0,
		nReduce: nReduce,
		nMap: len(files)}

	// Your code here.
	for idx, file := range files {
		c.toMapFiles[idx] = file
	}

	for i := 0; i < c.nMap; i++ {
		c.mapDone[i] = false
	}
	
	for i := 0; i < c.nReduce; i++ {
		c.toReduce[i] = false
		c.reduceDone[i] = false
	}

	c.server()
	return &c
}