advdupe2/lua/autorun/netstream.lua

387 lines
11 KiB
Lua

--A net extension which allows sending large streams of data without overflowing the reliable channel
--Keep it in lua/autorun so it will be shared between addons
AddCSLuaFile()
net.Stream = {}
net.Stream.SendSize = 20000 --This is the size of each packet to send
net.Stream.Timeout = 30 --How long to wait for client response before cleaning up
net.Stream.MaxWriteStreams = 1024 --The maximum number of write data items to store
net.Stream.MaxReadStreams = 128 --The maximum number of queued read data items to store
net.Stream.MaxChunks = 3200 --Maximum number of pieces the stream can send to the server. 64 MB
net.Stream.MaxSize = net.Stream.SendSize*net.Stream.MaxChunks
net.Stream.MaxTries = 3 --Maximum times the client may retry downloading the whole data
local WriteStreamQueue = {
__index = {
Add = function(self, stream)
local identifier = self.curidentifier
local startid = identifier
while self.queue[identifier] do
identifier = identifier % net.Stream.MaxWriteStreams + 1
if identifier == startid then
ErrorNoHalt("Netstream is full of WriteStreams!")
net.WriteUInt(0, 32)
return
end
end
self.curidentifier = identifier % net.Stream.MaxWriteStreams + 1
if next(self.queue)==nil then
self.activitytimeout = CurTime()+net.Stream.Timeout
timer.Create("netstream_queueclean", 5, 0, function() self:Clean() end)
end
self.queue[identifier] = stream
stream.identifier = identifier
return stream
end,
Write = function(self, ply)
local identifier = net.ReadUInt(32)
local chunkidx = net.ReadUInt(32)
local stream = self.queue[identifier]
--print("Got request", identifier, chunkidx, stream)
if stream then
if stream:Write(ply, chunkidx) then
self.activitytimeout = CurTime()+net.Stream.Timeout
stream.timeout = CurTime()+net.Stream.Timeout
end
else
-- Tell them the stream doesn't exist
net.Start("NetStreamRead")
net.WriteUInt(identifier, 32)
net.WriteUInt(0, 32)
if SERVER then net.Send(ply) else net.SendToServer() end
end
end,
Clean = function(self)
local t = CurTime()
for k, stream in pairs(self.queue) do
if (next(stream.clients)~=nil and t >= stream.timeout) or t >= self.activitytimeout then
stream:Remove()
self.queue[k] = nil
end
end
if next(self.queue)==nil then
timer.Remove("netstream_queueclean")
end
end,
},
__call = function(t)
return setmetatable({
activitytimeout = CurTime()+net.Stream.Timeout,
curidentifier = 1,
queue = {}
}, t)
end
}
setmetatable(WriteStreamQueue, WriteStreamQueue)
net.Stream.WriteStreams = WriteStreamQueue()
local ReadStreamQueue = {
__index = {
Add = function(self, stream)
local queue = self.queues[stream.player]
if #queue == net.Stream.MaxReadStreams then
ErrorNoHalt("Receiving too many ReadStream requests!")
return
end
for _, v in ipairs(queue) do
if v.identifier == stream.identifier then
ErrorNoHalt("Tried to start a new ReadStream for an already existing stream!")
return
end
end
queue[#queue+1] = stream
if #queue == 1 then
stream:Request()
end
return stream
end,
Remove = function(self, stream)
local queue = rawget(self.queues, stream.player)
if queue then
if stream == queue[1] then
table.remove(queue, 1)
local nextInQueue = queue[1]
if nextInQueue then
nextInQueue:Request()
else
self.queues[stream.player] = nil
end
else
for k, v in ipairs(queue) do
if v == stream then
table.remove(queue, k)
break
end
end
end
end
end,
Read = function(self, ply)
local identifier = net.ReadUInt(32)
local queue = rawget(self.queues, ply)
if queue and queue[1] then
queue[1]:Read(identifier)
end
end
},
__call = function(t)
return setmetatable({
queues = setmetatable({}, {__index = function(t,k) local r={} t[k]=r return r end})
}, t)
end
}
setmetatable(ReadStreamQueue, ReadStreamQueue)
net.Stream.ReadStreams = ReadStreamQueue()
local WritingDataItem = {
__index = {
Write = function(self, ply, chunkidx)
local client = self.clients[ply]
if client.finished then return false end
if chunkidx == #self.chunks+1 then self:Finished(ply) return true end
if client.downloads+#self.chunks-client.progress >= net.Stream.MaxTries * #self.chunks then self:Finished(ply) return false end
client.downloads = client.downloads + 1
local chunk = self.chunks[chunkidx]
if not chunk then return false end
client.progress = chunkidx
--print("Sending", "NetStreamRead", self.identifier, #chunk.data, chunkidx, chunk.crc)
net.Start("NetStreamRead")
net.WriteUInt(self.identifier, 32)
net.WriteUInt(#chunk.data, 32)
net.WriteUInt(chunkidx, 32)
net.WriteString(chunk.crc)
net.WriteData(chunk.data, #chunk.data)
if CLIENT then net.SendToServer() else net.Send(ply) end
return true
end,
Finished = function(self, ply)
self.clients[ply].finished = true
if self.callback then
local ok, err = xpcall(self.callback, debug.traceback, ply)
if not ok then ErrorNoHalt(err) end
end
end,
GetProgress = function(self, ply)
return self.clients[ply].progress / #self.chunks
end,
Remove = function(self)
local sendTo = {}
for ply, client in pairs(self.clients) do
if not client.finished then
client.finished = true
if CLIENT or ply:IsValid() then sendTo[#sendTo+1] = ply end
end
end
if next(sendTo)~=nil then
--print("Sending", "NetStreamRead", self.identifier, 0)
net.Start("NetStreamRead")
net.WriteUInt(self.identifier, 32)
net.WriteUInt(0, 32)
if SERVER then net.Send(sendTo) else net.SendToServer() end
end
end
},
__call = function(t, data, callback)
local chunks = {}
for i=1, math.ceil(#data / net.Stream.SendSize) do
local datachunk = string.sub(data, (i - 1) * net.Stream.SendSize + 1, i * net.Stream.SendSize)
chunks[i] = { data = datachunk, crc = util.CRC(datachunk) }
end
return setmetatable({
timeout = CurTime()+net.Stream.Timeout,
chunks = chunks,
callback = callback,
lasttouched = 0,
clients = setmetatable({},{__index = function(t,k)
local r = {
finished = false,
downloads = 0,
progress = 0,
} t[k]=r return r
end})
}, t)
end
}
setmetatable(WritingDataItem, WritingDataItem)
local ReadingDataItem = {
__index = {
Request = function(self)
if self.downloads+self.numchunks-#self.chunks >= net.Stream.MaxTries*self.numchunks then self:Remove() return end
self.downloads = self.downloads + 1
timer.Create("NetStreamReadTimeout" .. self.identifier, net.Stream.Timeout*0.5, 1, function() self:Request() end)
self:WriteRequest()
end,
WriteRequest = function(self)
--print("Requesting", self.identifier, #self.chunks)
net.Start("NetStreamWrite")
net.WriteUInt(self.identifier, 32)
net.WriteUInt(#self.chunks+1, 32)
if CLIENT then net.SendToServer() else net.Send(self.player) end
end,
Read = function(self, identifier)
if self.identifier ~= identifier then self:Request() return end
local size = net.ReadUInt(32)
if size == 0 then self:Remove() return end
local chunkidx = net.ReadUInt(32)
if chunkidx ~= #self.chunks+1 then self:Request() return end
local crc = net.ReadString()
local data = net.ReadData(size)
if crc ~= util.CRC(data) then self:Request() return end
self.chunks[chunkidx] = data
if #self.chunks == self.numchunks then self:Remove(true) return end
self:Request()
end,
GetProgress = function(self)
return #self.chunks/self.numchunks
end,
Remove = function(self, finished)
timer.Remove("NetStreamReadTimeout" .. self.identifier)
local data
if finished then
data = table.concat(self.chunks)
if self.compressed then
data = util.Decompress(data, net.Stream.MaxSize)
end
self:WriteRequest() -- Notify we finished
end
local ok, err = xpcall(self.callback, debug.traceback, data)
if not ok then ErrorNoHalt(err) end
net.Stream.ReadStreams:Remove(self)
end
},
__call = function(t, ply, callback, numchunks, identifier, compressed)
return setmetatable({
identifier = identifier,
chunks = {},
compressed = compressed,
numchunks = numchunks,
callback = callback,
player = ply,
downloads = 0
}, t)
end
}
setmetatable(ReadingDataItem, ReadingDataItem)
function net.WriteStream(data, callback, dontcompress)
if not isstring(data) then
error("bad argument #1 to 'WriteStream' (string expected, got " .. type(data) .. ")", 2)
end
if callback ~= nil and not isfunction(callback) then
error("bad argument #2 to 'WriteStream' (function expected, got " .. type(callback) .. ")", 2)
end
local compressed = not dontcompress
if compressed then
data = util.Compress(data) or ""
end
if #data == 0 then
net.WriteUInt(0, 32)
return
end
if #data > net.Stream.MaxSize then
ErrorNoHalt("net.WriteStream request is too large! ", #data/1048576, "MiB")
net.WriteUInt(0, 32)
return
end
local stream = net.Stream.WriteStreams:Add(WritingDataItem(data, callback, compressed))
if not stream then return end
--print("WriteStream", #stream.chunks, stream.identifier, compressed)
net.WriteUInt(#stream.chunks, 32)
net.WriteUInt(stream.identifier, 32)
net.WriteBool(compressed)
return stream
end
--If the receiver is a player then add it to a queue.
--If the receiver is the server then add it to a queue for each individual player
function net.ReadStream(ply, callback)
if CLIENT then
ply = NULL
else
if type(ply) ~= "Player" then
error("bad argument #1 to 'ReadStream' (Player expected, got " .. type(ply) .. ")", 2)
elseif not ply:IsValid() then
error("bad argument #1 to 'ReadStream' (Tried to use a NULL entity!)", 2)
end
end
if not isfunction(callback) then
error("bad argument #2 to 'ReadStream' (function expected, got " .. type(callback) .. ")", 2)
end
local numchunks = net.ReadUInt(32)
if numchunks == nil then
return
elseif numchunks == 0 then
local ok, err = xpcall(callback, debug.traceback, "")
if not ok then ErrorNoHalt(err) end
return
end
local identifier = net.ReadUInt(32)
local compressed = net.ReadBool()
if numchunks > net.Stream.MaxChunks then
ErrorNoHalt("ReadStream requests from ", ply, " is too large! ", numchunks * net.Stream.SendSize / 1048576, "MiB")
return
end
--print("ReadStream", numchunks, identifier, compressed)
return net.Stream.ReadStreams:Add(ReadingDataItem(ply, callback, numchunks, identifier, compressed))
end
if SERVER then
util.AddNetworkString("NetStreamWrite")
util.AddNetworkString("NetStreamRead")
end
--Send requested stream data
net.Receive("NetStreamWrite", function(len, ply)
net.Stream.WriteStreams:Write(ply or NULL)
end)
--Download the sent stream data
net.Receive("NetStreamRead", function(len, ply)
net.Stream.ReadStreams:Read(ply or NULL)
end)